Partially Observable Reinforcement Learning with Memory Traces

Thank you very much for your positive review! Below, we respond to your critical points and questions.

Expanding the discussion on connections to prior work on POMDPs. We agree that our discussion on related works is relatively sparse. Based on your suggestions, as well as the other reviewers' suggestions, we will expand this section with a discussion on connections to existing empirical work that focuses on using RNNs and transformers for reinforcement learning in POMDPs.

Why is the reward signal a function of the observations and not of the states? These two descriptions of POMDPs are equivalent. Given a POMDP with state space $\mathcal X$, observation space $\mathcal Y$, set of possible rewards $\mathcal R$, emission probabilities $p(y \mid x)$, and state-reward function $r: \mathcal X \to \mathcal R$, we can consider the equivalent POMDP with observation space $\tilde{\mathcal Y} = \mathcal Y \times \mathcal R$, emission probabilities $p(\tilde y = (y, r) \mid x) = p(y \mid x)[r = r(x)]$, and observation-reward function $\tilde r(\tilde y = (y, r)) = r$. We use the observation-reward definition because it emphasizes that rewards are also quantities that are observed by the agent. In this model, state and reward are conditionally independent given the observation. With the state-reward definition this is not the case, so rewards would technically need to be included explicitly in memory as they could also give information about the state. We are open to changing our definition if it is confusing.

Why is the policy Markovian? If the policy $\pi$ is Markovian, then the combination of POMDP and $\pi$ describe a hidden Markov model with the same state space $\mathcal X$ as the POMDP. If the policy is not Markovian, then this is no longer true. However, this is irrelevant to our paper, and so we will change the definition in the final version as follows. Consider the general setup consisting of a POMDP with state space $\mathcal X$, observation space $\mathcal Y$, action space $\mathcal U$, transition dynamics $p(x_{t + 1}\mid x_t, u_t)$ and emission probabilities $p(y_t\mid x_t)$, and a general agent with agent state space $\mathcal Z$, agent state update probabilities $\pi(z_{t+1}\mid z_t, y_{t+1})$ and policy $\pi(u_t\mid z_t)$. Together, these describe a hidden Markov model with state space $\Psi = \mathcal X \times \mathcal Z$, observation space $\mathcal Y$, transition dynamics $p(\psi_{t+1} = (x_{t + 1}, z_t) \mid \psi_{t} = (x_{t}, z_{t-1})) = \sum_{y_t}p(y_t\mid x_t)\pi(z_t\mid z_{t - 1}, y_t)\sum_{u_t}\pi(u_t\mid x_t)p(x_{t + 1}\mid x_t, u_t),$ and emission probabilities $p(y_t\mid \psi_t = (x_t, z_{t-1})) = p(y_t\mid x_t).$ Thus, as none of our results depend on the state space of the HMM, this definition shows that our results also apply with non-Markovian policies.

Are there any structural assumptions on the POMDP for which one or the other is preferable instead? This is an interesting question that we do not yet know the answer to. Our theoretical results suggest that it is always possible to use memory traces and perform at least as well as the best window-based approach. However, this is only true for the optimal choice of $\lambda$, which in general will not be known. We are interested in investigating if there is a way to connect properties of the transition and emission dynamics of the POMDP to the optimal choice for $\lambda$.

If $\lambda > 1/2$, are there environments where windows outperform memory traces? Yes, there are (e.g., see example A.1). However, there is a "simple fix": just choose a different $\lambda$. If $\lambda$ is set to a value less than $1/2$, Theorem 5.5 guarantees that no window will outperform the memory trace. The T-maze, which we analyze in Theorem 5.8, is qualitatively different: there is no "simple fix" to make a window-based approach efficient here (e.g., by choosing a different $m$). Thus, there is no environment where a window-based approach outperforms memory traces for all $\lambda$. However, there are environments where a memory trace approach outperforms windows for all $m$ (e.g., the T-maze).

Other suggestions. Thank you very much for your additional suggestions. We are working on a follow-up project right now, and will have a look at your references on RL with trajectory-dependent objectives. We are also examining if it is possible to learn $\lambda$ online.

Thank you very much for your positive review! Below, we respond to your critical points and questions.

Larger experimental domains would make the paper more attractive.

Indeed, the environments we consider in our experiments are more illustrative than realistic. We intend the main contributions of this paper to be theoretical, and we are working on a follow-up project in which we investigate memory traces empirically in more complex environments. The reason for this is that a more thorough empirical evaluation requires a lot more work. For example, the assumption of one-hot observation only makes sense in small tabular POMDPs. A proper treatment of these issues therefore goes beyond the scope of this paper.

It appears that the effectiveness of the method relative to the windows method depends on the value of the forgetting factor. What determines the forgetting factor itself? Under what conditions do we need a forgetting factor > 1/2?

These are interesting questions that we do not yet know the answers to. All that our theoretical (and empirical) results suggest is that $\lambda$ should usually be larger than $\frac{1}{2}$, as otherwise a window could be used instead. In our follow-up work, we investigate whether it is possible to learn $\lambda$, and if there is a way to connect properties of the transition and emission dynamics of the POMDP to the optimal choice for $\lambda$.

Thank you very much for your positive review! Below, we respond to your critical points and questions.

I think it would have been interesting to discuss the use of RNNs and transformers to represent the history of observations.

We are happy to include a brief discussion of the connections to RNN and transformer-based methods. Indeed, it is possible to view memory traces as a kind of simplified RNN, while the window-based approach has a fixed context length like transformers. Based on the other reviews, we will include references to https://arxiv.org/abs/1507.06527 and https://arxiv.org/abs/2206.01078.

The evaluation is on relatively simple tasks. Exploring a broader range of environments could further establish the generality of the approach.

Indeed, the environments we consider in our experiments are more illustrative than realistic. We intend the main contributions of this paper to be theoretical, and we are working on a follow-up project in which we investigate memory traces empirically in more complex environments. The reason for this is that a more thorough empirical evaluation requires a lot more work. For example, the assumption of one-hot observation only makes sense in small tabular POMDPs. A proper treatment of these issues therefore goeses beyond the scope of this paper.

It would be interesting to discuss how to choose $\lambda$ in different environments.

We agree that this is an interesting question that is not fully addressed in the paper. All that our theoretical (and empirical) results suggest is that $\lambda$ should usually be larger than $\frac{1}{2}$, as otherwise a window could be used instead. In our follow-up work, we investigate whether it is possible to learn $\lambda$, and if there is a way to connect properties of the transition and emission dynamics of the POMDP to the optimal choice for $\lambda$.

Do you assume the agent remembers its actions? Actions also provide important context for the agent.

In this paper, we define memory traces only based on the observations. However, it is easy to include actions in the traces as well by regarding them as observations, i.e., by considering the augmented observation space $\tilde{\mathcal Y} = \mathcal Y \times \mathcal U$.

I am more familiar with different descriptions of the POMDP setting, where the reward is a mapping from the latent state to the reward. Does this impact the proposed method?

These two descriptions of POMDPs are equivalent, so it would not impact the proposed method. Given a POMDP with state space $\mathcal X$, observation space $\mathcal Y$, set of possible rewards $\mathcal R$, emission probabilities $p(y \mid x)$, and state-reward function $r: \mathcal X \to \mathcal R$, we can consider the equivalent POMDP with observation space $\tilde{\mathcal Y} = \mathcal Y \times \mathcal R$, emission probabilities $p(\tilde y = (y, r) \mid x) = p(y \mid x)[r = r(x)]$, and observation-reward function $\tilde r(\tilde y = (y, r)) = r$. We use the observation-reward definition because it emphasizes that rewards are also quantities that are observed by the agent. In this model, state and reward are conditionally independent given the observation. With the state-reward definition this is not the case, so rewards would technically need to be included explicitly in memory as they could also give information about the state. We are open to changing our definition if it is confusing.

Thank you very much for your thorough review of our paper! Below, we respond to your critical points and questions. Due to limited space, we cannot respond to every point, but we will do our best to incorporate all your suggestions in the final version.

Proof of Theorem 5.8 and the choice of $\lambda = \frac{k - 1}{k}$ in the T-maze. The full proof of Theorem 5.8 is in Appendix B (p. 14). In this proof, as well as in our PPO experiment, we fix the forgetting factor to $\lambda = \frac{k - 1}{k}$, where $k$ is the corridor length of the T-maze. Importantly, this corridor length is fixed, and does not change during an episode; it is not to be confused with the time step $t$. Thus, $\lambda$ is not dynamic, and the memory trace is indeed an exponentially weighted average.

Concerns about the PPO (T-maze) experiment

• In the PPO experiment, we additionally give the memory trace agent access to the current observation (or, equivalently, we use a second memory trace with $\lambda = 0$). We did not do this to give an edge over the frame stacking agent. Indeed, the frame stacking agent also has access to the current observation. Instead, it was a heuristic choice and a demonstration that a practitioner is not constrained to using just one memory trace. We are happy to run additional experiments with only a single memory trace.
• We use the Adam optimizer with a linear learning rate decay from $0.0003$ to $0$ as suggested by [Huang et al. 2022].
• In the T-maze minigrid, the agent receives a reward of $1$ if it succeeds, and a reward of $0$ if it does not (and $0$ in all nonterminal transitions). Thus, the "average success rate" is the same as the "average total reward". A random agent would achieve an average success rate of $0.5$. The observed performance slightly below $0.5$ is due to episodes exceeding the maximum episode length of $5(k + 2)^2$ steps in a T-maze with corridor length $k$. We will clarify this in the article. Does this address your concerns about the Y-axis in Figure 5?

Concerns about the TD (random walk) experiment

• Thank you for your feedback. In this experiment, our goal was to compare memory traces to the full window approach whose sample efficiency we analyze in Sec. 5. However, we agree that we could instead concatenate observations to improve sample efficiency, as we do in the PPO experiments. While we are not aware of this approach in the context of linear function approximation, we ran the experiments that you suggested, please see here. All three algorithms ran for 100,000 steps (1/100th of Fig. 4) with a grid search over step sizes. We see that concatenation makes it possible to learn with very long windows, and that this approach almost reaches the best value error of memory traces. However, it requires almost 20 times the memory to achieve this, and it lacks the theoretical guarantees that we develop in Sec. 5.
• The step size used in the original TD experiments (Fig. 4) for the window approach is $\alpha = 0.0001, 0.0002, 0.005, 0.05, 0.2, 0.5$ for $m = 1, 2, 3, 4, 5, 6$, respectively. These values were chosen by a grid search optimizing the average value error. For the memory trace, the step size was kept constant at $\alpha = 0.0002$ for all values of $\lambda$. This value was also determined by a grid search.

Further points

• We will expand Appendix C in the final paper. (optimizer, step sizes, ..)
• Essential References. Thank you for these suggestions. We will include them in the final version.
• Relation to eligibility traces. We do believe that there is a close connection to eligibility traces, and will include an extended discussion in the final version.
• Definition of sample efficiency. A low minimum value error does not correspond to low sample efficiency. Could you please explain what gave this impression? We use Theorem 5.1 to tie sample efficiency to the metric entropy.
• Definition of environment $\mathcal E$. In our definition of $\mathcal E$, the policy is assumed fixed (line 107).
• More realistic experiments & observation-reward function. Please see our response to Reviewer tJLm.
• Q1: Lack of injectivity guarantee. Could you please elaborate on this question? We provide one experiment in a continuing environment and one in an episodic environment.
• Q2: "Traces cannot outperform windows". In a fully observable POMDP, we can use $\lambda = 0$ to achieve the same result. Theorem 5.5 guarantees that, if learning with windows is tractable, then learning with traces is also tractable.

Please let us know if you still have any concerns, or if our responses were not satisfactory. If you agree that we have addressed your questions and provided the necessary clarifications, we kindly request raising our paper’s score based on the responses provided. We are happy to incorporate further feedback into the paper to improve clarity and remove ambiguity.