Predicting Future Actions of Reinforcement Learning Agents

We are thankful for the reviewer’s comments and for recognizing the significance of the problem statement. The questions are addressed as follows:

What's the rationality behind the choices in different methods?

The rationality is mainly based on our intuition of which information is useful. We agree that more comprehensive experiments would benefit the paper. As such, we conducted further experiments to show the impacts of choosing different types of inner-state. Please refer to our global response for the results. We will include these results in the revised paper.

How do you handle the potential divergence between simulated trajectories and real trajectories over time?

First, we note that our problem statement does not permit access to real trajectories, so predictions must be made purely on the basis of simulated trajectories and the current state, which may diverge depending on the quality of the world model. We make no special effort to prevent this divergence, although we note that the world model is trained on transitions the agent encounters during training.

To cope with the potential divergence when predicting future actions and events, we input the simulated trajectories to a neural network, as the neural network should learn to extract useful information from the trajectories while discarding the non-useful information. In the extreme case of the world model outputting random noise, such that the simulated trajectories contain almost no information, the neural network can still rely on the current state to make a prediction (which is equivalent to the baseline case shown in Fig. 2). We also note that it is due to this possible divergence that we do not just use the empirical distribution of simulated trajectories to manually predict future actions and events.

How are the training/testing samples generated? Are they from trajectories with fixed length?

We deploy the trained agents on the training level of Sokoban until 50k transitions are collected. Each episode may last from a few steps to over a hundred steps, depending on whether and how fast the trained agent can solve the level. We perform the same process for generating testing samples, but we deploy the trained agent on testing levels instead of training levels and only collected 10k transitions.

We are thankful for the reviewer’s suggestions and hope that our response adequately addresses the reviewer’s questions. We hope that the reviewer considers revising the score in light of our new experiments and clarification.

We are thankful for the reviewer’s comments and for recognizing the significance of the problem statement. The questions are addressed as follows:

W1: As the goal of the figure is to compare the effectiveness of different approaches on different agents instead of the prediction variability, it is indeed better to use a 2-sigma error bar that indicates the confidence interval instead. We will update the paper to use 2-sigma error bar. Please see the updated figures in our global response for the revised figures.

W2: We designed the event to be predicted to have a small probability, so it is closer to real-world applications where we may want to predict an unsafe but rare event. In our trained agents, the agent will step on the blue tile within L=5 steps around 5% of the time. Therefore, a trivial predictor that predicts the future event not occurring will already achieve 95% accuracy. As such, we report F1 score instead of accuracy in event prediction. Such a trivial predictor will only achieve a zero F1 score due to zero recall, making the problem far from trivial. We expect that an even rarer event will make achieving a high F1 score even more challenging due to the scarcity of the data.

We are thankful for the reviewer’s suggestions and hope that our response adequately addresses the reviewer’s questions. We hope that the reviewer considers revising the score in light of our clarifications.

We are thankful for the reviewer’s detailed comments. The questions are addressed as follows:

W1: Please see our global response. Briefly, our results largely accord with our hypotheses, and we will add the hypotheses to the introduction following the list of research questions.

W2: We define planning as “the process of interacting with an environment simulator or a world model to inform the selection of subsequent actions” (lines 71-72). Our distinction is meant to capture the difference between agents that do this explicitly by construction (explicit planning agent), those that behaviorally seem to do so implicitly by learning (implicit planning agent), and those that behaviorally do not seem to do so at all (non-planning agent). The exact definition of these three types of agents can be found in lines 75-96.

The type of planning is important when considering predicting future actions or events as different types of information are available. With an explicit planning agent, the rollouts in the world model provide an easily interpretable plan (e.g., we can directly look at the rollouts in Fig. 1 of the Thinker paper). With an implicit planning agent, the hidden activations may contain the same degree of information as the explicit rollouts in explicit planning agents. However, this information is distributed across neurons and likely mixed with other internal computations that are not relevant to its plans, making it harder to extract the agent's plans. Lastly, for a non-planning agent, the lack of planning means that the hidden state may only contain a more compact representation of the state. It is expected that these different agents' inner states have different degrees of information and intractability. This expectation is largely consistent with the results shown in Fig. 2.

W3: In the simulation-based approach, we unroll a given agent in a trained world model instead of the environment. Therefore, it is expected that the trained world model predicts the same state space as the input to the agent. In special cases where the action can be computed using only the abstract state (e.g., the actor in Dreamer only requires the latent state to compute the action distribution), we can indeed use an abstract world model (e.g., the same latent world model in Dreamer) to do the simulation. We are thankful that the reviewer pointed this out and will clarify it in the revised paper.

P.S. Latent predicted state may be less effective for event prediction compared to the raw predicted state. We briefly experimented with a latent world model using Thinker and found that the event prediction accuracy dropped significantly, likely because the information in latent representation is harder to extract. However, if we train a decoder to convert the latent predicted state to the raw predicted state, the accuracy restores to the level reported in the paper. As such, we advise decoding the latent predicted state into the raw predicted state even when rollouts in latent space are possible.

W4: We will move Section 5 after the introduction to improve clarity. As aforementioned, the paper considers a new problem setting, so we are unable to find any related work that tackles the same problem. As such, we can only list the works done in related areas and/or share similar goals, such as predicting opponent actions in a multi-agent setting or preventing (instead of only detecting) dangerous agent actions in a safe RL setting.

W5: Please see our answer to the three research questions in our global response. In particular, we advise using the simulation-based approach when the environment dynamics are easy to model, and using the inner-state approach in other cases. We also advise using explicit planning agents if the inner-state approach is considered, as their explicit plans are generally more informative or easily extractable than the inner states of other agents.

Why is section 4 titled Methods? There are no method contributions in this work.

We propose two methods, and both methods are part of the paper’s contributions. As this is a new problem setting, the two methods are constructed to be general and simple to allow future work to build upon them.

We note that the simulation-based approach has been used in a multi-agent setting to predict opponent movements for training a better policy, but the nature of the problem is very different from the one we consider. For the inner-state approach, we are not aware of any similar methods used in other problem settings.

Shift line 159-160 before line 145.

Thank you. We will update the paper accordingly.

Why do authors write that environments with diverse states do not exist?

The diversity here refers to the variety of possible solutions (or optimal policies) as in Sokoban, where there are usually multiple ways of solving the level, leading to high difficulty in predicting agent actions. We did not posit that such environments do not exist, and welcome pointers to alternative environments with this property, as we have not been able to identify any. We briefly experimented with Atari environments, where in some games the state space is very diverse, but we found that the lack of diverse optimal policies reduces the problem of future action prediction to predicting the optimal actions.

What is meant by 'false alarms' in line 350?

It means that when the predictor forecasts that an agent is going to perform an unsafe action, but in fact, the agent won’t be doing it. This may lead to improper responses (such as shutting down the agent) that are not warranted. We will rephrase the broader statement to make it clearer.

We are thankful for the reviewer’s detailed feedback. We will update the paper as stated above and hope that the reviewer considers revising the score in light of the clarified hypothesis and context. We are also happy to discuss further if the reviewer has any additional questions.

Thank you for the response and the feedback.

I am still not convinced that the hypotheses and the empirical results go beyond anything currently not known about the different types of planning agents, particularly the ones considered in the paper. Errors in policy learning in 'explicit' planning agents based on an incorrectly learned world model is already well understood. Similarly, the results for 'implicit' planning agents are not surprising either. A more objective analysis, even for a single category of approaches, could make the paper much stronger.

We respectfully disagree with the assertion that the hypothesis and empirical results are already known. To the best of our knowledge, there has not been any study on the predictability of different RL agents. We would appreciate any references to relevant papers if they exist. Our experimental results partially confirm our hypothesis, while some findings are relatively more surprising (as stated in the global response). It’s important to note that the error in policy learning is distinct from what we are considering here—we are investigating the predictability of RL agents, not the process of learning a policy.

I do not agree with the authors' definition of planning.

We did not propose any new definitions for planning or a new taxonomy of RL agents in the paper. Our definition of planning is consistent with previous work, such as [1] and [2]. There are many proposed definitions of planning. We do not believe there is a universally agreed-upon definition.

Recent works [3, 4] suggest that agents can also learn to perform planning with the correct architecture. Therefore, we believe that further distinguishing between explicit and implicit planning agents is beneficial for discussing the results, as labeling them as planning-free contradicts the findings of these works. This distinguishment between explicit and implicit planning agents is not novel, as it has already been addressed in the literature—see, for example, Table 5 in [5], which distinguishes between explicit and implicit planning. Developing a rigorous mathematical definition for explicit and implicit planning agents is far from trivial (and was also not attempted in [5]), as it is difficult to determine whether an agent learns to plan within its neural network. This challenge is well beyond the scope of our paper.

In summary, our classification of agents follows existing work, and we have not introduced any new definitions or made any claims to contribute to the taxonomy of existing RL agents. We will reference [5] in the revised paper when discussing the agent types to clarify that this is not a new taxonomy. We included the taxonomy solely to aid in discussing the methods and results.

[1] Chung, S., Anokhin, I., & Krueger, D. (2024). Thinker: learning to plan and act. Advances in Neural Information Processing Systems, 36.

[2] Hamrick, J. B., Friesen, A. L., Behbahani, F., Guez, A., Viola, F., Witherspoon, S., ... & Weber, T. On the role of planning in model-based deep reinforcement learning. In International Conference on Learning Representations 2021.

[3] Guez, A., Mirza, M., Gregor, K., Kabra, R., Racanière, S., Weber, T., ... & Lillicrap, T. (2019, May). An investigation of model-free planning. In International Conference on Machine Learning (pp. 2464-2473). PMLR.

[4] Garriga-Alonso, A., Taufeeque, M., & Gleave, A. (2024). Planning behavior in a recurrent neural network that plays Sokoban. In ICML 2024 Workshop on Mechanistic Interpretability.

[5] Moerland, T. M., Broekens, J., Plaat, A., & Jonker, C. M. (2023). Model-based reinforcement learning: A survey. Foundations and Trends® in Machine Learning, 16(1), 1-118.

It is hard to interpret the proposed taxonomy as 'Methods' in the current version of the work. A proposed method is a formal framework and/or approach, which is not a contribution in the current version of this work.

Section 4 proposes using inner-state and simulated rollouts to aid in predicting future actions and events. We introduced new types of additional information that can enhance prediction and experimentally demonstrated that they indeed improve prediction accuracy. A taxonomy typically involves a simple categorization of existing methods, but we have emphasized that both of these methods are novel. Simultaneously proposing two methods in a new problem setting does not constitute creating a taxonomy. We strongly disagree with the reviewer’s assessment that our two proposed methods should be dismissed as merely a taxonomy.

We are thankful for the reviewer’s comments and for recognizing the significance of the problem statement. The questions are addressed as follows:

If the action space is rolled out and appended to the input, would this not result in the potential future actions already being included with the input?

These rollouts are rollouts of imaginary actions performed in a learned world model and thus do not constitute data leakage. This allows action and event predictions to be made L=5 steps before the true future action or event is known. The future action may indeed be included in simulated rollouts of model-based agents, but these agents (including Thinker) won't necessarily choose the actions that are simulated in the rollouts due to numerous reasons, such as an agent learning to ignore bad rollouts, adjusting for model prediction error, or continuously refining the plan in future steps. Nonetheless, the simulated actions in rollouts are certainly helpful for predicting future actions—this is expected, and this potential advantage of explicit planning algorithms is one of the core motivations of our work.

The reviewer is also unsure as to the benefits of the method over, for example, training a world model, rolling out several series of actions and choosing the branch that accumulates the highest predicted reward to also serve as a predictor of future actions.

We interpret the suggestion as using the empirical action or event distribution of the rollout with the highest rollout returns (sum of predicted rewards and the terminal value) as the predicted future action and event distribution.

For model-based RL agents such as Thinker and MuZero, they do not necessarily select the rollout with the highest rollout return. Figure 19 in the appendix of the Thinker paper shows that a trained Thinker agent selects the next action with the highest rollout return around 90% of the time. In contrast, the same figure for MuZero (shown as MCTS) is around 50%, as MCTS selects the next action based on visit count instead*.

For model-free RL agents, as they cannot perform rollouts like Thinker or MuZero, their selected actions likely differ from the best action from multiple rollouts. They often exhibit myopic behavior, such as pushing a box to a wall and thus making the level unsolvable. The empirical rollouts (i.e., rollouts using the agent policy) should be more insightful for predicting future actions and events.

As such, if we had to estimate the future action or event prediction by empirical distribution of a rollout, we suggest using**:

1. The rollout with the highest return for Thinker.
2. The rollout with the largest visit count for MuZero.
3. The empirical rollout for model-free RL agents.

However, in our paper, we consider inputting these rollouts to a neural network to learn the prediction instead of directly using the empirical distribution. This is because the world model may diverge from the true environment, leading to unreliable rollouts. For example, in the extreme case that the trained world model outputs random noise, the learned neural network can still predict future actions or events based on the current state while ignoring the rollout, resulting in the baseline case shown in Figure 2. This should be considerably better than using the empirical distribution of rollouts, which does not provide any information in this hypothetical example.

Another advantage of using a learned neural network is that we only need the binary event label instead of the event function when training the predictor. On the other hand, if we use the empirical distribution of rollouts to make event predictions, we also need to handcraft a function that maps predicted states to an event label (e.g., classify whether the agent stands in the blue location from raw predicted pixels in our setting).

*Note that both figures will be lower if we consider the action in L steps (which we are interested in predicting) instead of just the next action.

**The first two are included in the inner-state approach, while the third is included in the simulation-based approach.

We hope that the above answer clarifies the reviewer’s questions and look forward to further discussion.