Paper: Revisiting Familiar Places in an Infinite World: Continuing RL in Unbounded State Spaces (ICLR-2024)

Reviewer comment (Novelty):
There exists a line of work on using training wheels for RL in real systems, e.g., [Towards Safe Online Reinforcement Learning in Computer Systems](http://mlforsystems.org/assets/papers/neurips2019/towards_mao_2019.pdf), that this work ignores.

Author response:
Thank you for pointing out this work. We will include it in the related work sections. The work you cite is actually a deep RL version of the work by [1] that we already cite in our paper. Our comments in the related work section for [1] also apply to this new work you mention: “they assume that *a stable policy is given* and use it as a starting point for learning an optimal policy. We make no such assumption and try to learn a stable policy directly from a random policy”. We would expect methods that have access to a stable policy to perform competitively, but the main point is that we cannot always rely on access to a stable policy, and in this work, we explore the setting without access to this fixed stable policy.

Reviewer comment (Theory):
How can state get unbounded in a real problem? It seems to me that there is always going to be a physical limit to the state space of real systems, e.g., maximum queue lengths in the queuing environment that you explored.

Author response:
We understand this point. We acknowledge this in the footnote on page 1. Moreover, in the traffic control description in Section 5.1, we say: “ In this domain, SUMO models a real-life traffic situation by imposing a cap on the number of cars per lane. However, as we will show, existing deep RL algorithms still fail despite this cap”. Thus, the challenges in the unbounded state space arise even in the bounded state space setting where the bound is high.

In general, a bounded state space with a high bound is not substantially different from an unbounded state space, especially in multi-dimensional problems. An unstable policy will diverge towards the boundary of the state space instead of towards infinity; see Figure 2 (d)-(f). In contrast, our approach can quickly learn to maintain the system in low-cost states and away from the boundary.

Reviewer comment (Experiments):
As explained in the second paragraph of 5.2., the choice of cost function and its interplay with stability objective is problem specific. How should one choose a good cost function?

Author response:
Yes, we acknowledge this point. We do want to point out that a linear stability cost $\ell(s) = \|s\|$ is sufficient for achieving stability in most of our environments. We need to use the reciprocal stability cost only in a highly stochastic queue environment. Also, using a simple quadratic stability cost $\|s\|^2$ (inspired by domain knowledge from control theoretic and queueing problems) is sufficient in ALL our environments.

We can think of adjusting the cost function as introducing domain knowledge, which is common in many RL approaches for real-world tasks. In our case, and as mentioned in Section 5.2, we made the adjustment based on relative magnitude scales between the optimality and stability cost.

Reviewer comment (Theory):
The paper seems to emphasize the average cost may be unbounded and difficult to optimize, e.g., in Section 4.1 "As we described in Section 3, $J^O$ is difficult to minimize directly; even keeping $J^O$ bounded is challenging". However, in this setting, generally we are not directly optimizing this objective but optimizing a differential return (also mentioned in the paper). The differential value function is bounded, and trust region or PPO style algorithm can be applied.

Author response:
Yes PPO is learning the differential value function.  However, the differential value function is well-defined only when the policy is already stable with a bounded average cost and a well defined stationary distribution. Algorithmically, average-reward PPO and TRPO methods need to first estimate the average cost before estimating the differential values; for example, see Line 4 in Algorithm 2 of https://proceedings.mlr.press/v139/zhang21q/zhang21q.pdf

Reviewer comment (Reproducibility):
The paper lacks an algorithm description for their own method. Though several changes are described in Section 4, it is unclear to readers what value functions are used for shaped reward and original reward, how the policy is updated.

Author response:
We have included pseudo-code in the updated appendix (section H)

Reviewer comment (Presentation):
This paper considers the setting of average reward. It is better to use a subsection named "average reward MDPs" to give an overview. Though some basic terms are included in Section 2, a well organization is helpful for readers not familiar with this area.

Author response:
We understand. We will include the necessary background in the Appendix.

Reviewer comment (Theory):
What is the intuition of Assumption 3 in Section 2? Why is divergence in the cost equivalent to the divergence in the state?

Author response:
Thank you for the prompt reply!

While there are only two actions available to the agent, the difficulty lies in poor credit assignment to these two actions due to high stochasticity, poor reward functions, and unbounded state spaces. While both actions move the agent towards the origin, the agent may not actually be able to move towards the goal since with some probability two things can happen (both of which can happen at the same time-step): 1) its movement can fail (stays in the same place) and 2) it may be pushed away from the goal. Thus, while both actions are "good" in the sense they will try to make the agent move towards the goal, the actual quality of an action is determined based on the environment stochasticity.

Suppose the Manhattan distance from the origin is $c$ and is large (so the agent incurs a cost of $c$). Consider the agent is in a given fixed state and is deciding between the two actions.

1) Suppose the agent takes action $a_1$ and reduces the distance to the goal by at most 1 (since it can move only one step towards the goal and environment stochasticity can push the agent away by two steps). Thus the new cost incurred is $c'_1 = c - 1$. 

2) if the agent takes $a_2$ and it increases the distance by at most 1 (for the similar reason as before),  the new cost incurred is $c'_2 = c + 1$. 

If $c$ is large, then the costs incurred in both these situations is $\approx c$ (since $c >> 1$). At that point, the agent finds it difficult to distinguish between the quality of these two actions since it incurs the similar cost for taking either action. The inability to distinguish between good and bad actions results in divergence. We refer you to Section 3 and Section C in the Appendix as well.

Of course, if the true values for the policy were actually learned, the agent could distinguish between long-term good and bad actions. However, to get the true values the agent needs many samples of the above transition due to the high stochasticity. Our paper is showing that in our setting optimizing the optimality cost functions and unbounded state spaces (see Section 3) in the reset-free and highly stochastic setting makes convergence to the true values challenging, making it difficult for the agent to determine good/bad actions, which makes the agent diverge.

To mitigate this challenge, our paper is proposing the agent first optimize a stability cost (and then optimality) instead of directly optimizing the optimality cost since the former can be more informative to the agent.

Hopefully this helps clarify things!

Reviewer comment (Theory):
In section 4.3, authors mention the reward shaping and scheduling can guide agent to "re-visit regions of states with low costs". I am not very sure about this claim. Since reward shaping is dealing with a new cost $c(s_t) - c(s_{t-1})$, why learning policy on this cost can ensure agent to visit state $s$ with low cost $c(s)$?

Author response:
We thank the reviewer for responding and acknowledging our response!
 
Yes, minimizing the stability cost $c(s_t) – c(s_{t-1})$ encourages the agent to transition to states with lower norm or same norm (incurs negative or zero cost). Once the agent is stable, to ensure that the agent *actually* goes to states with low true cost $c(s_t)$, we must then include the optimality term ($c(s_t)$) which is the full version of our STOP algorithm.
 
Thank you for your question! We hope our answers and updated pdf have resolved your concerns and you would consider re-evaluating your score.

Reviewer comment (Theory):
While the state space could be unbounded since it cannot be controlled in most cases, the cost function is usually manually designed. If an unbounded reward function is problematic, a natural solution would be to design a better reward function, which appears to be an easier fix.

Author response:
**Re: suggested reward functions.**

In Section 3, we show that with such bounded reward functions the agent still diverges. In the updated Appendix (section F), we also include the performance of an RL agent using the suggested reward functions and show that the agent still diverges.

Part of our message is that designing reward functions for this class of problems is difficult. Common techniques of bounding and normalizing (such as in Section 3 and the suggested reward functions in the review) are not easily applicable in this setting. Thus, we need to resort to other techniques to get RL agents to perform well in these settings.

We also emphasize that in many stochastic control problems, the original cost function is tied to a specific performance metric. For example, in queueing, minimizing the average queue length is equivalent to minimizing system latency. Changing the cost function will in general change the optimal policy (i.e., the optimal policy of the new MDP no longer minimizes latency). Our reward shaping approach has the property that the optimal policy remains unchanged, since the average stability cost $\lim_{T\to\infty} (1/T) \sum_{t=1}^T \mathbb{E}[g(s_{t-1}, s_t)]$ is 0 for any stable policy (and hence for the optimal policy as well); see Proposition 2.

Reviewer comment (Experiments):
For the tasks mentioned in this work, the hardness of the tasks could be mainly due to poorly designed reward functions. For example, can you compare methods with the reward function \( r(s,a,s') = \exp(-||s+a||_2^2) \) or \( r(s,a,s') = \exp(-||s'||_2^2) \) in the gridworld task?

Author response:
**Re: suggested reward functions.**

In Section 3, we show that with such bounded reward functions the agent still diverges. In the updated Appendix (section F), we also include the performance of an RL agent using the suggested reward functions and show that the agent still diverges.

Part of our message is that designing reward functions for this class of problems is difficult. Common techniques of bounding and normalizing (such as in Section 3 and the suggested reward functions in the review) are not easily applicable in this setting. Thus, we need to resort to other techniques to get RL agents to perform well in these settings.

We also emphasize that in many stochastic control problems, the original cost function is tied to a specific performance metric. For example, in queueing, minimizing the average queue length is equivalent to minimizing system latency. Changing the cost function will in general change the optimal policy (i.e., the optimal policy of the new MDP no longer minimizes latency). Our reward shaping approach has the property that the optimal policy remains unchanged, since the average stability cost $\lim_{T\to\infty} (1/T) \sum_{t=1}^T \mathbb{E}[g(s_{t-1}, s_t)]$ is 0 for any stable policy (and hence for the optimal policy as well); see Proposition 2.

Reviewer comment (Theory):
Section 4.3 state is poorly motivated. Why the state transformations can solve the extrapolate burden? I have no idea why just scaling the input to the network and make it "feels like a visited state" can improve anything. The claim "all these functions reduce divergence rate" is not grounded.

Author response:
**Re: visiting states and divergence comment.**

The state transformations do not solve the extrapolation burden. They *only mitigate* it. For example, suppose in the 2-queue setting, we have queue lengths [100, 200] and the RL agent has learned the optimal actions for this state. Now if the agent observes queue lengths [2000, 3000], which are very far from [100, 200], it will have a difficult time determining the optimal actions since it has to perform *extreme* generalization, which neural networks are bad at performing [1]. However, if we use a symlog transformation, as done in the paper, then the agent has learned optimal actions for symlog([100, 200]) = [4.6, 5.2], and can then *locally* generalize [1] to symlog([2000, 3000]) = [6.9, 7.6], which is significantly closer to the states that the agent was trained on.

Please see our updated Appendix (section D and E) and global response. In Section E, we show that when states are transformed using symlog, the agent observes the transformed states a higher fraction of the time compared to observing the original states.

Reviewer comment (Theory):
The theoretical results is not deep enough to motivate the paper. For example I have no idea the what the proposition 1 and 2 is doing. Yes we have a policy gradient on this \( g \) function and yes now the \( J^S \) is rate stable. So what do these mean? Why proving these two propositions can prove the proposed method is better in any aspect?

Author response:
**Re: theoretical result**

The theoretical results prove the soundness of our proposed method. 
1. Proposition 2  shows that if an agent optimizes the stability cost, it is indeed stable i.e. it will incur bounded cost, which means it will visit states with bounded norm.
2. Proposition 1 shows that the stability cost can be optimized by policy gradient methods, as it has a well-defined gradient *even when the current policy is unstable*.

This should be contrasted with optimizing the optimality cost alone  (without the stability cost). When initialized from a random policy (which is likely to be unstable), it will not have a finite cost or policy gradient, so the agent cannot even start running a policy gradient method.

Reviewer comment (Presentation):
The presentation of the experiment is confusing. I almost consider the SUMO experiment is omitted in the paper as the Figure 2 mixes result from 3 environments without the introduction and the discussion of the evaluation metrics. We should have a section (either in Appendix or main body and has clear annotation) to discuss the very detailed meaning of the evaluation metrics and when presenting in the experiment section we need to emphasize what metric in what environment is improved/affected by what.

Author response:
**Re: experiment section**

Thanks for the suggestion. We do include the evaluation metrics and environment descriptions in the Appendix (section B), but we will improve this presentation.

Reviewer comment (Presentation):
It will be good to unify the color scheme in Fig 2, like setting "O"'s color to green across all figures.

Author response:
**Re: color coding.**

Sure, we can do this for the camera ready.

Reviewer comment (Theory):
Eq 3 looks similar to the Lagrangian reward shaping method in Safe RL. Can we setting the cost as the primary objective and set the stability cost to be the constraint? We can write Eq 3 to be like \( c(s_t)  - \lambda g(s_{t-1}, s_t) \) and use traditional RL to solve. So that we can avoid the adhoc design of the annealing scheme.

Author response:
**Re: stability as constraint**

We agree that our approach bears similarity to a Lagrangian formulation. Empirically, we observe that using a fixed lambda, even when chosen by hyperparameter sweeping, still leads to divergence. In the updated Appendix (section G), we hyperparameter sweep over different lambda and use your proposed formulation ($c(s_t) + \lambda g(s_{t-1}, s_t)$; note that the constraint is added instead of subtracted when learning) and show that the agent diverges. We note that if we combine your technique with the bounded cost/reward suggestion from reviewer Nors, the unbounded nature of the state will make the optimality reward close to 0, which will effectively make the agent optimize the stability cost/reward only, which is essentially our STOP agent without the optimality component.

Treating and updating lambda as a Lagrangian multiplier is an interesting idea. Note that unlike usual constrained problems where the objective and constraint are competing, in our problem stability (the constraint) is a necessary first step to optimality (the objective). If the current policy is not yet stable, the first term in the Lagragian formulation $c(s_t) + \lambda g(s_{t-1}, s_t)$ will be diverging and hence dominating the objective. It is therefore crucial to first achieve stability and maintain stability, as done in STOP. Also, ensuring stability requires a constraint of the form $\lim_{T\to\infty} (1/T) \sum_{t=1}^T \mathbb{E}[g(s_{t-1}, s_{t})] = 0$, which is a hard, equality, average-cost constraint, making it challenging to directly apply existing constrained RL methods.

Reviewer comment (Theory):
Why we use the world cost instead of reward. I guess it's because in this infinite horizon setting our goal is to reduce the "average-cost" like the Eq 1.

Author response:
**Re: use of cost instead of reward.**

As mentioned in Section 2, these terms are from the control-theoretic perspective. Given our applications (queueing, traffic control) are typically researched from the control-theoretic perspective, we stuck to that terminology.

Reviewer comment (Experiments):
We are now using \( g = c(s_{t-1}) - c(s_t) \). Would that possible to apply the idea of multi-step TD like \( g = c(s_{t-10}) - c(s_t) \) or even with the idea of \( TD(\lambda) \)? What would happen in this case.

Author response:
**Re: multi-step**

This is an interesting suggestion, and is definitely something we will consider in future work. We expect that performance would be problem-specific, and as is typical with multi-step methods, will depend on the hyperparameter $\lambda$

Reviewer comment (Experiments):
The annealing is too adhoc. Can we have a mechanism to dynamically adjust the preference over cost & stability cost by observing the state uncertainty, either via a distributional Q values or via some kinds of representation learning approaches?

Author response:
**Re: dynamically setting annealing schedule**

This is another very interesting direction. Our main point of the work is to show that an agent should first pursue stability and then optimality. To get this point across, we resorted to a fixed schedule used by [2]. Dynamically setting the weights is in itself a difficult research problem, but we agree a starting point would be to explore methods that consider state uncertainty.