Paper: Using Forwards-Backwards Models to Approximate MDP Homomorphisms (ICLR-2024)

Reviewer comment (Theory):
The assumptions are a bit limiting – deterministic dynamics and state-dependent rewards.

Author response:
**Regarding limiting assumptions**

We agree deterministic dynamics is limiting, which is why we also validate our approach in stochastic environments (pred-prey and Asterix). Taking your advice onboard, we have further addressed this limitation by adding more MinAtar games.

For the state dependent rewards, we agree that this is a limitation of our work. However, our approach is still applicable to a range of discrete control environments (e.g. all the MinAtar games). If there is an important class of environments we have not listed in our limitations, please let us know so that we can include it.

Reviewer comment (Theory):
One major weakness is that I don’t know the connection between the forward and backward model (Eq 4 and 5), and MDP homomorphism (Eq 1 and 2). Why not directly learn Eq 1 and 2 (like my question below)?

Author response:
**Regarding your proposal for learning a homomorphism by directly optimizing on the definition of an MDP homomorphism with an MSE loss**

Thank you for this suggestion but, as far as we are aware, this would not work. Only optimizing the MSE loss based on the MDP homomorphism would mean there is nothing to prevent the homomorphic state mapping from learning a trivial representation (e.g. mapping all states to a vector full of zeros would give zero loss). Van der Pol et al. (2020a) back up this claim explicitly (see page 3 left column of their paper) and get around this issue by using a contrastive loss (which we compare to in the Cartpole and Predator Prey experiments, however, there has not been any attempt to scale their dynamic programming approach to larger inputs, which is why we do not compare to it in MinAtar).  Rezaei-Shoshtari et al. (2022) get around this by adding an approximate bisimulation metric to their loss function in addition to the MDP homomorphism loss function. We think it could be possible to balance these two terms in the discrete setting but would be a significant project in its own right.

Reviewer comment (Writing):
I don’t understand Eq 6. What is sigma(s,a), not sigma(s)? No definition. Also the hypothetical action part.

Author response:
- Apologies for not being clear on this, now in our homomorphism definition we describe how our abstract states have the same dimensionality as they are predicted as states that you would see from your environment. This means the LHS and RHS can be equal in Eq. 6.

Reviewer comment (Writing):
The writing quality is fair. Notation writing can be improved. Be complete on $\mathcal R(s,a)$ and $\mathcal P(s’\mid s,a)$ in the background. $G$ should be in expectation.

Author response:
**Regarding clarity and mistakes in the writing**

- return is now in expectation
- $\sigma$ is now always a function of $(s, a)$, as our mapping is action dependent
- removed the hat in the algorithm
- figure 3(a) typo is fixed 
- starting sentence with \forall has been fixed
- double which on page 3 fixed
- we have fixed the issues with \citep and \citet

all of these changes are in the updated manuscript.

Reviewer comment (Writing):
Define the symbols like \sigma and \alpha_s explicitly with inputs and outputs (you can check how van der Pol et al., 2020b define them).

Author response:
- we provide textual description and explicit input, output function signatures for $\sigma$ and $\alpha$
- our paper is now consistent in that all $\sigma$ are written as $\sigma_{a}$ to note the dependent on actions, this is different from the Van der Pol definition and we make note of it. Is that satisfactory?

Reviewer comment (Experiments):
The experiments are only conducted on 4-5 tasks, which are not enough for evaluating an algorithm (as the experiments are the bulk of this work).

Author response:
Thank you for taking the time to review our work. **One major criticism is that we need more experimental verification. We agreed with this criticism. We have now run all the MinAtar games with the model now being learned online with a new model-based baseline also included.**

Reviewer comment (Experiments):
Asterix results are not in favor of the algorithm, and more MinAtar environments are expected to run for a comprehensive comparison.

Author response:
Thank you for taking the time to review our work. **One major criticism is that we need more experimental verification. We agreed with this criticism. We have now run all the MinAtar games with the model now being learned online with a new model-based baseline also included.**

Reviewer comment (Experiments):
For me, a straightforward approach to learn homomorphism in discrete action space, is to parameterize value function as $Q(\sigma(s), \phi_s(a))$, then train it end-to-end with some Q-learning loss, and also train $\sigma$, $\phi_s$ with MSE losses to satisfy (1)(2). It is similar to the approach in the continuous action version of homomorphism (Rezaei-Shoshtari et al., 2022). The action can be one-hot encoded. I am a bit surprised that this simple approach is not discussed in the paper, and what do you think of this approach?

Author response:
**Regarding your proposal for learning a homomorphism by directly optimizing on the definition of an MDP homomorphism with an MSE loss**

Thank you for this suggestion but, as far as we are aware, this would not work. Only optimizing the MSE loss based on the MDP homomorphism would mean there is nothing to prevent the homomorphic state mapping from learning a trivial representation (e.g. mapping all states to a vector full of zeros would give zero loss). Van der Pol et al. (2020a) back up this claim explicitly (see page 3 left column of their paper) and get around this issue by using a contrastive loss (which we compare to in the Cartpole and Predator Prey experiments, however, there has not been any attempt to scale their dynamic programming approach to larger inputs, which is why we do not compare to it in MinAtar).  Rezaei-Shoshtari et al. (2022) get around this by adding an approximate bisimulation metric to their loss function in addition to the MDP homomorphism loss function. We think it could be possible to balance these two terms in the discrete setting but would be a significant project in its own right.

Reviewer comment (Evaluation):
The return on Asterix is much lower than reported by prior papers. The original nature DQN paper has return 6012 on Asterix, but why is the return reported by this paper below 20? Is there any rescaling of the reward?

Author response:
**Regarding Asterix Returns**

It is true that the returns are different to the normal Atari Asterix returns. This is because we are using MinAtar and there is not a one to one mapping between the returns in Atari (Bellemare et al. 2013) and MinAtar (Young and Tian 2019).

Reviewer comment (Writing):
The notations are not clear and sometimes confusing. For example, the hypothetical actions are sometimes referred to as $a_{hyp}$ and sometimes as $\hat{a}$. In Eq (7), $a^\prime$ is never defined. $\sigma$ is sometimes a function of states, sometimes a function of state-action pairs. In Eq (8), should the first $s^\prime$ be $\bar{s^\prime}$? Although the meanings can be roughly inferred from the context, it is still not reading-friendly.

Author response:
**Regarding lack of clarity in notation**

We thank you for pointing out some lack of clarity in our notation. We have addressed it in the following ways:

- $\sigma$ is now always defined as a function of states and actions.
- we now always refer the hypothetical action as $a_{hyp}$
- we define a’ in our definition of the MDP homomorphism
- we clarify in the text that in equation 8, $s’=\bar{s}’$ because both transitions lead to the same next state

All of these changes are in the updated manuscript.

Reviewer comment (Experiments):
Is there any experimental result with learned forward/backward model on large-scale environments?

Author response:
**Regarding more MinAtar results with learned models**

The criticism that sticks out to us most seems to be that we do not have sufficient results where the transition models are learned online. As a result, we have rewritten the MinAtar sections with new results where the transition model is learned online for all MinAtar games for a fair comparison with the baselines. 

We found different optimizers work best for different algorithms and environments. To take this effect into account, we compare the ratio of performance improvement from EEA to a standard DQN for each optimizer and average over optimizers (RMS and Adam). We also average the performance improvement over all games. This is the result shown in Figure 4 in the updated manuscript. As you can see, overall, EEA delivers a nice performance boost in the low sample limit but its asymptotic performance is slightly worse after a large number of frames.

Reviewer comment (Theory):
Formalism lacks the precision needed to understand the main technical concepts. For instance, a "forwards model" and "backwards model" are never defined. These seem like parameterized functions. However, their input-output signatures are not described, and their parameters seem to change from Equation 6 to Algorithm 1.

Author response:
- We now provide more details on the input and output signatures of our models after equation 5.  
- We now describe explicitly that all models are trained with mean squared error loss and experience tuples from the DQN experience replay buffer  
the parameters change from algorithm 1 to equation 6 is a typo we have fixed

Reviewer comment (Experiments):
Current methodology fails to account for the pre-training experience used to learn the EEA.

Author response:
- We have run our approach on all MinAtar games where we learn the model online (and so account for model learning experience)

Reviewer comment (Experiments):
Experiments are missing important baselines. For instance, a standard model-based learner that doesn't also learn a homomorphism.

Author response:
- We add a model-based Dyna baseline to compare to for MinAtar

Reviewer comment (Theory):
The main contribution is potentially limited, as it only applies to fully-observable, deterministic MDPs where rewards are strict functions of state (i.e., no action dependence).

Author response:
While it is true our approach is defined theoretically for deterministic MDPs, we test what happens when we break this assumption by running on stochastic environments on predator prey and the MinAtar environments.

Reviewer comment (Writing):
Mathematical notation is sometimes overloaded and often undefined.

Author response:
- We now provide more details on the input and output signatures of our models after equation 5.  
- We now describe explicitly that all models are trained with mean squared error loss and experience tuples from the DQN experience replay buffer  
the parameters change from algorithm 1 to equation 6 is a typo we have fixed

Reviewer comment (Writing):
Many references are missing parentheses. This makes some sentences difficult to read.

Author response:
- We have addressed the issues with parentheses in our citations

Reviewer comment (Theory):
This assumption is not generally true: "Equivalent effect abstraction is based on a simple observation that is guaranteed to be true in the majority of RL problems if R and P are deterministic."

Author response:
Lastly, if possible, could you provide some more detail or example benchmark environments where our assumptions about value equivalence are not true? That way we can understand your criticism on this point better.

Reviewer comment (Experiments):
The domains selected for experiments seem well-suited to equivalent effect abstraction because they are generally spatial "navigation" tasks in 2D - I'm not sure how to describe the domains precisely, but where hypothetical actions are easy to select and equivalent effects are prominent.

Author response:
**“My main concern is the applicability of this approach beyond the domains that have been tested to showcase it … I'm not sure how to describe the domains precisely, but where hypothetical actions are easy to select and equivalent effects are prominent. It would be interesting to see this idea applied in the other MinAtar games.”**

We agree that this is a useful next step. As a result, we have rewritten the MinAtar section of our paper with new results on all games with a model that is learned online. We find that the benefit is most obvious in games like Asterix and Seaquest, where the number of actions collapsed down to the hypothetical action is large. Nevertheless, we also see a nice improvement on games like freeway when using the RMS optimizer, where the hypothetical action is not easy to select.

Reviewer comment (Experiments):
Why was asterix chosen as the only benchmark tested from MinAtar?

Author response:
**Why was asterix chosen as the only benchmark tested from MinAtar?**

Asterix was initially chosen as a benchmark because it has a large number of equivalent actions, so hence is a good demonstration of our approach. Since our initial submission we have run EEA on all of the Minatar games. We hope this addresses some of your concerns about the general applicability of our approach. We applied EEA to some games where the action structure isn’t as simple (e.g Seaquest which has a “fire” action or Freeway where the speed of moving forward depends on previous actions). While some of these results are positive, we do agree that our approach is very well suited to navigation tasks.