Paper: FP-IRL: Fokker-Planck-based Inverse Reinforcement Learning --- A Physics-Constrained Approach to Markov Decision Processes (ICLR-2024)

Reviewer comment (Novelty):
“Most IRL methods require the transition function to be known”: This is repeatedly claimed in the paper and is considered as the motivation behind this paper. This is not true as a number of works propose sample-based estimators to solve IRL [1,2,3,4,5,6,7]. These methods should be discussed and compared against.

Author response:
- To reiterate our motivation as discussed in the 2nd and 3rd paragraph in Section 1 and last paragraph in Section 3.1:
    > This is motivated by real-world examples (e.g. cancer cell metastasis), where we can only collect the trajectory data but dynamic knowledge is missing or imperfect, and therefore accessibility of sampling from transition is not available.
Many systems (e.g., swarms, crowd behavior) have mechanistic foundations in physics, which if exploited can lead to a better understanding and more efficient learning of their incentive structures (i.e., reward function) and transition. With the above motivation, we propose a new method of physics-constrained IRL. This method simultaneously estimates the transition rooted in physics and reward functions using only data on trajectories, while also inferring physical principles that govern the system and using them to constrain the learning.
- Our method is different than other IRL methods as they would require either the accessibility to sampling from the transition or the analytical form of the transition. Our method does not require both. To reiterate our problem definition: Our target problem setting is that both reward and transition are unknown and need to be inferred in the MDP (i.e., $\mathcal{M} / \{ T(\cdot), R(\cdot) \}$).
Specifically, the transition is unknown in any of these forms:
    - The analytical form of the transition function $T(s' | s, a)$ is unknown.
    - It is not accessible for sampling the next state for a given state and action (i.e., $s' \sim T(s' | s, a)$ cannot be realized). Please note that conventional IRL approaches often assume that sampling is available even though they claim their method does not require the transition (i.e., they do not require the analytical form of the transition but require the ability to sample from the transition.). This may lead to confusion for the audience.
- We did not conduct an extensive comparison experimentally to other IRL methods because most existing IRL methods require the transition function to be known or **the ability to sample the next state from the transition function**. This is not consistent with our problem setting as these methods do not provide a systematic way to infer the transition or build a model for sampling. It's not directly comparable as the method we propose is designed to simultaneously estimate the transition and rewards under physics constraints.
- Discussion on given papers: (1) None of these papers assume the transition is not accessible for sampling, and provide a method to infer the transition function as we did. (2) Therefore, they all require at least the sampling accessibility from the transition function in the reward inference (e.g., for collecting and comparing the trajectory of learned policy, or constructing the replay buffer for policy optimization using learned reward). A detailed discussion of each paper is provided below: 
    - Ref. [1]: In lines 2 and 5 of Algorithm 1, the transition is required for the estimation of expected state frequencies of learned policy. 
    - Ref. [2]: In the objective function Eq. (7) on page 4, the analytical form of transition (i.e., transition matrix) is required.
    - Ref. [3]: In line 4 of Algorithm 1, accessibility to sampling from the transition is required for collecting trajectories and policy optimization (this paper uses TRPO).
    - Ref. [4]: In Section 5.1, it requires the sampling from the transition for constructing the replay buffer. 
    % Moreover, this method uses IRL for imitation learning. Although it claims it does not require explicit transition knowledge.  
    - Ref. [5]: It requires sampling from the transition either for online or offline settings. 
    % imitation learning, replay buffer
    - Ref. [6]: In line 1 of Algorithm 1, it requires sampling from the transition to construct the replay buffer.
    - Ref. [7]: In step 1 in the for-loop of Algorithm 1, it requires sampling from the transition to construct the replay buffer and collecting trajectories of learned policy.

Reviewer comment (Experiments):
“Empirical treatment of dynamics using NN’s is not generalizable”: While this statement is true, the way this problem is addressed is not by FP but by using Hermite cubic functions since FP also requires estimating the potential function which has the same learning complexity under an expressible function class. How do other prior methods compare if you use the Hermite cubic function as limited hypothesis class.

Author response:
- As discussed in Sections 1 and 5, we mean that the empirically estimated transition using data alone may not generalize well to the state and action regions away from the data, and may not follow the underlying dynamics.
- In this work, processes whose time-continuous form is governed by the FP equation are of fundamental interest. Therefore, by first inferring the governing FP equation (using VSI) and using Conjecture 3.1, we have a transition function that follows the underlying dynamics and has an unambiguous interpretation in terms of physics. The Hermite polynomial basis is only for the convenience of representation and by itself does not grant generalization. This problem indeed is addressed by the FP constraint on the MDP dynamics.
- DL and Finite Element approximations offer different ways of parameterizing functions. We choose the Hermite basis functions (a finite element basis) in contrast to DL as it offers a systematic way of approximating differentiable functions. A deep network would not have the issue of discretizing the FP equation by the FE method and bring more scalability in higher dimensional problems. However, the DNN would not allow us to separately extract the potential function from the gradient fields unless a specially designed modular architecture is used for the DNN. Additionally, the expressivity that any chosen DNN depends on the network architecture i.e. to approximate a given function with arbitrary precision, one may require really large networks. This scenario is similar to $h^p$-convergence in mesh-based methods, where higher precision is achieved by reducing the mesh sizes. In our experience, the mesh-based method offers better control in terms of defining the functional spaces used to approximate the value function. Moreover, the tradeoff is that it would be less physically interpretable and more disconnected from the governing physics, and it would require larger trajectory datasets. So under situations where interpretability is less important and data is ample, DL could be an appropriate choice.

Reviewer comment (Theory):
“This means that the agent employs to reach states whose value function is high”: This is not a true statement as some high value states can be unreachable under a given starting state distribution. A better explanation for the conjecture might be needed.

Author response:
- In the last sentence of Section 3.3 on page 5, we have the condition that "every state in the MDP is reachable". Our method specializes in problems that follow Fokker Planck dynamics at a continuum scale, so this would not be a limitation.

Reviewer comment (Theory):
This paper hinges on a conjecture, which might be problematic. This conjecture says that the $\psi$ is equivalent to the $Q$ function in RL problem. However, this $\psi$ function is related to the dynamics. In particular, $\nabla \psi$ is the gradient field. It seems unclear to me why the conjecture is true because the $Q$ function depends on the choice of reward function as well. It is possible that the authors actually assume that the equivalence is between the (entropy-regularized) optimal policy and the (entropy-regularized) optimal Q function. In this case, it is still unclear to me why the conjecture should be true. It would be great if the authors could at least prove it in special cases such as linear-quadratic case.

Author response:
Thank you for your reviews.
### Reply to Weakness 1 and Question 1:
- We agree that there remains theoretical work to be carried out to provide principled guarantees for Conjecture 3.1. There also remains considerable scope for empirical evidence. We hope to work in this direction in future work. For now, we note that this work is motivated by problems of cancer cell dynamics for which our collaborators have generated the real-world example that we used. Cell migration dynamics are widely understood to be governed in the continuous limit by different versions of FP equations. This is the class of problems that this work addresses, by basing itself on Conjecture 3.1. Furthermore, many other real-world problems including Brownian dynamics, swarming and crowd behavior, pattern formation, and morphogenesis are described by FP equations in the continuous limit.  Understanding their incentive structure (i.e., reward function) holds great potential for understanding these complex systems. 
We are not claiming that the conjecture holds uniformly for MDPs. That is likely not the case.

Reviewer comment (Novelty):
For discrete-time MDPs, there are existing works that study max-entropy IRL without knowing the transition function. In fact, not knowing the reward seems not a barrier -- we can just estimate the transition from the trajectory data. Thus, the claim that "While most IRL approaches require the transition function to be prescribed or learned a-priori, we present a new IRL method targeting the class MDPs that follow the It^{o} dynamics without this requirement." seems ungrounded. Moreover, it seems unclear to me why Ito dynamics should be used here.

Author response:
### Reply to Weakness 2:
- Our method is motivated by real-world problems (e.g. cancer cell problems) whose dynamic is known to be governed by physics, specifically that described by the Ito dynamics. A detailed reason for using Ito dynamics can be seen in our motivation as discussed in the 2nd and 3rd paragraph in Section 1 and last paragraph in Section 3.1:
    > This is motivated by real-world examples (e.g. cancer cell metastasis), where we can only collect the trajectory data but dynamic knowledge is missing or imperfect, and therefore accessibility of sampling from transition is not available.
Many systems (e.g., swarms, crowd behavior) have mechanistic foundations in physics, which if exploited can lead to a better understanding and more efficient learning of their incentive structures (i.e., reward function) and transition. With the above motivation, we propose a new method of physics-constrained IRL. This method simultaneously estimates the transition rooted in physics and reward functions using only data on trajectories, while also inferring physical principles that govern the system and using them to constrain the learning.
- Our method is different than other IRL methods as they would require either the accessibility to sampling from the transition or the analytical form of the transition. Our method does not require both. To reiterate our problem definition: Our target problem setting is that both reward and transition are unknown and need to be inferred in the MDP (i.e., $\mathcal{M} / \{ T(\cdot), R(\cdot) \}$).
Specifically, the transition is unknown in any of these forms:
    - The analytical form of the transition function $T(s' | s, a)$ is unknown.
    - It is not accessible for sampling the next state for a given state and action (i.e., $s' \sim T(s' | s, a)$ cannot be realized). Please note that conventional IRL approaches often assume that sampling is available even though they claim their method does not require the transition (i.e., they do not require the analytical form of the transition but require the ability to sample from the transition.). This may lead to confusion for the audience.
- > not knowing the reward seems not a barrier" 
    - If the reviewer means "not knowing the 'reward' seems not a barrier": Yes, the goal of the IRL problem is to infer the reward function from the trajectory data.
    - If the reviewer means "not knowing the 'transition' seems not a barrier": there are two ways to infer the transition in the IRL problem:
       - Inferring the transition before the reward inference. This is what this reviewer means. As we discussed in Section 5 and in 2nd paragraph in Section 2: 
            > "The empirically estimated transition may not inherit the underlying dynamics, and is often challenging to generalize to state and action regions away from training samples relying on data alone."

            > "The lack of interpretability in deep learning models (inferred before the reward inference) can translate to difficulty in scientific understanding of the system behavior.'"
        - Inferring the transition simultaneously with reward inference. This is what we try to do. 
            > "By constraining with FP dynamics, we systematically reduce this ambiguity to identify a unique pair of transition and reward functions that comply with the FP dynamics." 
            
            This is more beneficial and preferable when the system has a mechanistic foundation and we can exploit it into the inference to reduce the ill-posedness and add physical interpretability.

Reviewer comment (Novelty):
Related work. This work is also related to the huge literature on solving inverse problems involving PDEs, for example, neural ODE and PINN.

Author response:
### Reply to Weakness 3:
- Please note that VSI is a published method [Wang et al. (2019; 2021)] that provides a way to infer FP PDE. Our major contribution in this work is FP-IRL algorithm.
- PINNS or neural ODE also offers a way to infer the PDE, but they are out of the scope of this IRL work and not included in the related work.

Reviewer comment (Experiments):
The evaluation/experiments is quite "weak". There are no baselines and no comparison to state of the art. It is important to know where the new method stands in terms of performance against other state of the art methods with no physics priors.

Author response:
Thank you for your reviews.

### Reply to Weakness 1 and Question 1:

- Our target problem setting is that both reward and transition are unknown and need to be inferred in the MDP (i.e., $\mathcal{M} / \{ T(\cdot), R(\cdot) \}$).
Specifically, the transition is unknown in any of these forms:
    - The analytical form of the transition function $T(s' | s, a)$ is unknown.
    - It is not accessible for sampling the next state for a given state and action (i.e., $s' \sim T(s' | s, a)$ cannot be realized). Please note that conventional IRL approaches often assume that sampling is available even though they claim their method does not require the transition (i.e., they do not require the analytical form of the transition but require the ability to sample from the transition.). This may lead to confusion for the audience.
- We did not conduct an extensive comparison experimentally to other IRL methods because most existing IRL methods require the transition function to be known or **the ability to sample the next state from the transition function**. This is not consistent with our problem setting as these methods do not provide a systematic way to infer the transition or build a model for sampling. It's not directly comparable as the method we propose is designed to simultaneously estimate the transition and rewards under physics constraints.

Reviewer comment (Evaluation):
No timings are provided for the experiments. Although I tend to agree with the statement "FP-IRL avoids such iterations altogether and instead induces a regression problem leveraging the FP physics that is also computationally more stable", wall-time performance can have very different outcomes (e.g. training the model in the proposed method takes too long).

Author response:
### Reply to Weakness 2 and Question 2:
- We regret that we did not measure the computation time of our method when we conducted the initial experiment on the cluster, but we reproduced the experiment using lower discretization resolution on the laptop with the Apple M1 Max chip and 32 GB memory: 

For the synthetic toy example in Section 4.1 (Please note that original discretization resolution ranges from 5 to 17 in the convergence analysis):
|     Number of Discretization |            8 |             9 |           10 |           11 |
|------------------------------|-------------:|--------------:|-------------:|-------------:|
|     FEA Time (sec)           |        77.05 |         128.6 |       253.56 |       344.04 |
|     Regression Time (sec)    |         0.77 |          1.12 |         1.64 |         2.46 |
|     Total VSI Time (sec)         |        77.82 |        129.72 |        255.20 |        346.50 |

For the modified Mountain Car problem in Appendix F.2 (Please note that original discretization resolution ranges from 10 to 30 in the convergence analysis):
| Number of Discretization |   20  |   25  |   30   |
|:------------------------:|:-----:|:-----:|:------:|
| FEA Time (sec)           | 45.95 | 88.54 | 156.29 |
| Regression Time (sec)    | 0.87  | 1.71  | 3.17   |
| Total VSI Time (sec)         | 46.82 | 90.25 | 159.46 |

Reviewer comment (Reproducibility):
More details about the VSI method (in the main text) could help an ML-oriented reader.

Author response:
### Reply to Weakness 3:
- VSI details are provided in Appendix E and cited papers [Wang et al. (2019; 2021)] as indicated in the second sentence of Section 3.6. 
- Please note that VSI is a published method that provides a way to infer FP PDE. Our major contribution in this work, as discussed in Section 1, is the FP-IRL algorithm.

Reviewer comment (Experiments):
Can you provide some baselines of state of the art methods with no physics priors?

Author response:
Thank you for your reviews.

### Reply to Weakness 1 and Question 1:

- Our target problem setting is that both reward and transition are unknown and need to be inferred in the MDP (i.e., $\mathcal{M} / \{ T(\cdot), R(\cdot) \}$).
Specifically, the transition is unknown in any of these forms:
    - The analytical form of the transition function $T(s' | s, a)$ is unknown.
    - It is not accessible for sampling the next state for a given state and action (i.e., $s' \sim T(s' | s, a)$ cannot be realized). Please note that conventional IRL approaches often assume that sampling is available even though they claim their method does not require the transition (i.e., they do not require the analytical form of the transition but require the ability to sample from the transition.). This may lead to confusion for the audience.
- We did not conduct an extensive comparison experimentally to other IRL methods because most existing IRL methods require the transition function to be known or **the ability to sample the next state from the transition function**. This is not consistent with our problem setting as these methods do not provide a systematic way to infer the transition or build a model for sampling. It's not directly comparable as the method we propose is designed to simultaneously estimate the transition and rewards under physics constraints.

Reviewer comment (Evaluation):
Can you provide timings of the proposed method?

Author response:
### Reply to Weakness 2 and Question 2:
- We regret that we did not measure the computation time of our method when we conducted the initial experiment on the cluster, but we reproduced the experiment using lower discretization resolution on the laptop with the Apple M1 Max chip and 32 GB memory: 

For the synthetic toy example in Section 4.1 (Please note that original discretization resolution ranges from 5 to 17 in the convergence analysis):
|     Number of Discretization |            8 |             9 |           10 |           11 |
|------------------------------|-------------:|--------------:|-------------:|-------------:|
|     FEA Time (sec)           |        77.05 |         128.6 |       253.56 |       344.04 |
|     Regression Time (sec)    |         0.77 |          1.12 |         1.64 |         2.46 |
|     Total VSI Time (sec)         |        77.82 |        129.72 |        255.20 |        346.50 |

For the modified Mountain Car problem in Appendix F.2 (Please note that original discretization resolution ranges from 10 to 30 in the convergence analysis):
| Number of Discretization |   20  |   25  |   30   |
|:------------------------:|:-----:|:-----:|:------:|
| FEA Time (sec)           | 45.95 | 88.54 | 156.29 |
| Regression Time (sec)    | 0.87  | 1.71  | 3.17   |
| Total VSI Time (sec)         | 46.82 | 90.25 | 159.46 |

Reviewer comment (Experiments):
Not broadly applicable to more challenging problems, since discretization is required. As the state-action space increases, the approach would become infeasible.

Author response:
- VSI can fundamentally handle high-dimensional state-action spaces. However, having its root in finite element methods, this method does suffer from the curse of dimensionality. More precisely, considering a $d$-dimensional state-action space, with a uniform grid of $n$-bins, we get a discrete state-action space with $n^d$ points. The number of operations in the numerical setup of VSI scales linearly with these points i.e. $\sim \mathcal{O}(n^d)$. Therefore, for a $k$-term model for the energy, we require $\sim \mathcal{O}(n^d k)$ operations that result in a linear regression problem with $k$-terms. These operations are vectorized and parallelizable whenever necessary.
- Please note that VSI is a published method [Wang et al. (2019; 2021)] that provides a way to infer FP PDE. Our major contribution in this work, as discussed in Section 1, is the FP-IRL algorithm, which inherently does not have the requirement on dimensionality. Future work includes modularizing the FP PDE inference: using other inference methods for higher dimension problems.
- As discussed in the **"Applicability"** part in Section 5, there are broad applications in scientific problems: 
    > "With physics-constrained modeling, this allows the application of FP-IRL to problems where the transition is not available for sampling and has not been mathematically modeled, or discovered. Cancer cell migration, as well as the migration of other cell types, is known to be governed by physics, specifically that described by the FP equation (Bressloff, 2014). Therefore, there is interest in the fields of biology, biophysics, and physics more broadly, to have scientific machine learning methods that respect these physics. We achieve this by combining machine learning ideas (IRL) with physics principles (Minimum Energy Principle and FP dynamics). Although the proposed method may not apply directly to some RL problem domains, such as robotics, many other physics phenomena encompassing Brownian dynamics (Keilson and Storer, 1952), swarming (Correll and Hamann, 2015) and crowd behavior (Dogbe, 2010), pattern formation and morphogenesis (Garikipati, 2017) are also described by FP equations in the continuous limit, and this work would also be applicable to them."

Reviewer comment (Experiments):
No comparison with other IRL methods was provided. For example, is there an advantage to using VSI for the cancer cell problem, over other simple IRL methods?

Author response:
- Our target problem setting is that both reward and transition are unknown and need to be inferred in the MDP (i.e., $\mathcal{M} / \{ T(\cdot), R(\cdot) \}$).
Specifically, the transition is unknown in any of these forms:
    - The analytical form of the transition function $T(s' | s, a)$ is unknown.
    - It is not accessible for sampling the next state for a given state and action (i.e., $s' \sim T(s' | s, a)$ cannot be realized). Please note that conventional IRL approaches often assume that sampling is available even though they claim their method does not require the transition (i.e., they do not require the analytical form of the transition but require the ability to sample from the transition.). This may lead to confusion for the audience.
- We did not conduct an extensive comparison experimentally to other IRL methods because most existing IRL methods require the transition function to be known or **the ability to sample the next state from the transition function**. This is not consistent with our problem setting as these methods do not provide a systematic way to infer the transition or build a model for sampling. It's not directly comparable as the method we propose is designed to simultaneously estimate the transition and rewards under physics constraints.
- Our motivation also demonstrates the advantage of FP-IRL in the cancer cell problem as discussed in the 2nd and 3rd paragraph in Section 1 and last paragraph in Section 3.1:
    > This is motivated by real-world examples (e.g. cancer cell metastasis), where we can only collect the trajectory data but dynamic knowledge is missing or imperfect, and therefore accessibility of sampling from transition is not available.
Many systems (e.g., swarms, crowd behavior) have mechanistic foundations in physics, which if exploited can lead to a better understanding and more efficient learning of their incentive structures (i.e., reward function) and transition. With the above motivation, we propose a new method of physics-constrained IRL. This method simultaneously estimates the transition rooted in physics and reward functions using only data on trajectories, while also inferring physical principles that govern the system and using them to constrain the learning.
- A general discussion of the advantages of FP-IRL over other methods is provided in the "Significance" part of Section 5.
- Again, please note that VSI is a published method that provides a way to infer FP PDE. Our major contribution is the FP-IRL algorithm.

Reviewer comment (Theory):
Page 5: "It imposes a physics constraint that the change in distribution should be small and approach zero over an infinitesimal time step". In the limit $\Delta t \rightarrow 0$, the free energy term should become negligible in Equation 9, right? Why can we then ignore the squared Wasserstein distance term instead of the free energy term?

Author response:
- The free energy $F(p, \psi)$ would not approach zero in an infinitesimal time step, but $\Delta t F(p, \psi)$ could as $\Delta t \rightarrow 0$. However, in an infinitesimal time step, the Wasserstein distance of two distributions $p_t(s, a)$ and $p_{t+1}(s, a)$ would be the dominant term vanishing to zero because it has higher order term on the difference term and change of distribution is small in an infinitesimal time step, and therefore can be neglected in the optimization problem. Then, the $\Delta t$ is a constant coefficient of the free energy and would affect the minimization.
- Note that Wasserstein-2 distance between two distributions $p_t(x)$ and $p_{t+1}(x)$:

$$ W_2^2(p_t, p) =  \inf_{\gamma \in \Gamma(p_t, p)} \int_{\mathcal{X} \times \mathcal{X}} \gamma(\boldsymbol{x}, \boldsymbol{y}) \left\| \, \boldsymbol{x}-\boldsymbol{y} \right\|_{2}^2 d\boldsymbol{x} d\boldsymbol{y} $$

where $\Gamma(p_t, p)$ is the set of coupling of $p_t$ and $p$. A coupling $\gamma$ is a joint probability measure on $\mathcal{X} \times \mathcal{X}$ whose marginals are $p_t $ and $p$ on the first and second factors, respectively.

Reviewer comment (Theory):
Does Equation 21 have a unique minimizer? If not, how does VSI address the unidentifiability/ill-posedness issue in IRL, i.e. there are many possible rewards (and corresponding value functions) that yield the same policy?

Author response:
- Equation 21 does not have a unique minimizer. Again, please note that VSI is a published method [Wang et al. (2019; 2021)] that provides a way to infer FP PDE. IRL ill-posedness is not reduced by VSI. 
- Instead, IRL ill-posedness is reduced by the physics constraints (i.e., FP dynamic) from our conjecture as we proposed in the paper. 
- As discussed in Section 5: 
    > By constraining with FP dynamics, we systematically reduce this ambiguity to identify a unique pair of transition and reward functions that comply with the FP dynamics.