Paper: Identifying Latent State Transition Processes for Individualized Reinforcement Learning (ICLR-2024)

Reviewer comment (Experiments):
The paper has a few explanation of why each part can work. It may be more helpful for readers to understand if the effect of each component and why it works can be more clearly displayed.

Author response:
**Q1: The paper has a few explanation of why each part can work. It may be more helpful for readers to understand if the effect of each component and why it works can be more clearly displayed.**

**A1**: Thanks for your constructive comments. We have carefully revised the manuscript to clearly explain the function and importance of each component. Specifically, we have improved Section 4 by adding the following:

1) An overview paragraph that outlines the design of our methodology and provides a brief summary of the overall framework.

> Our proposed framework is carefully customized to meet the requirements of the identifiability theorem. According to the identifiability definition A.3, latent confounders are identifiable if observational equivalence can consistently lead to the identifiability of the latent variables. This motivates our work to utilize a generative model and achieve confounder estimation by fitting the learned distribution to the true observed distribution. Traditional continuous VAE~\citep{kingma2013auto} encodes the input data into a continuous latent space using probabilistic encoders and then reconstructs the input from this space using decoders. However, this approach is not suitable for our problem setting. Specifically, given the offline trajectory of the population, we have customized the VAE framework and developed an individual-specific factor estimation framework to recover the latent factors.

> First, to obtain latent representations in a reduced-dimensional space, we use an encoder with the entire sequence $\vec{s}^k\_{0:T}$ of each individual $k$ as input and the latent continuous representation $z^k$ as output. Considering the group determinacy assumption, the estimated latent individual-specific factor should be discrete. We then incorporate a quantization layer to discretize the continuous latent representation $z^k$ into the latent individual-specific factor $\hat{l}\_k$. Intuitively, the better $\hat{l}\_k$ is recovered, the more effectively the data generation process can be learned. To reconstruct individualized transition processes, we design a conditional decoder that takes the conditions $(\vec{s}^k\_{t-1},\vec{a}^k\_{t-1})$, the estimated latent individual-specific factor $\hat{l}\_k$, and the noise term $\hat{\nu}\_k$ as inputs, producing $\hat{\vec{s}}^k\_t$ as the output. The noise is estimated based on $\hat{l}\_k$ to account for the stochasticity of the state transition processes. Afterward, the optimized quantized encoder is employed to augment the original buffer and facilitate downstream individualized policy learning.

2) A straightforward explanation of the functions of each component, including its inputs/outputs and essential mathematical formulations. Specifically,

-  In the quantized encoder module, we have added the motivation for introducing the quantization layer, along with corresponding update functions. The changes are as follows:

> The group determinacy assumption implies the existence of the latent individual-specific factor $L$. Since $L$ denotes the time-invariant latent individual-specific factors that influence each state in the transition process, we initially use an encoder to capture the continuous representation $z^k$ from $\vec{s}^k\_{0:T}$ and then employ a quantization layer to define unique and distinct groups as $\hat{l}^k$.

> The corresponding update functions are$o^k\_t=\text{Conv1D}(\vec{s}^k\_{t:t+H})$, $h^k\_t,c^k\_t=\text{LSTM}(h^k\_{t-1},c^k\_{t-1},\vec{s}^k\_t;\theta)$, where $o^k\_t$, $h^k\_t$ and $c^k\_t$ are the feature map, hidden state and cell state at time $t$, and $\theta$ denotes all the parameters of the LSTM.

-  In the noise estimation module, we have provided the motivation for including this module and explained how it applies within our framework.

> To account for the stochastic nature of the individualized transition processes, a noise estimator is introduced with $\hat{l}^k$ as input and $\hat{\nu}_k$ as output. 

> The motivations for introducing a noise layer are 1) to mimic the probabilistic property of the individualized state transition in Equation 1, and 2) to compensate for the loss of stochasticity due to the deterministic nature of the quantization operation.

> Our posterior $ q(\hat{l}\_k=\kappa|\vec{s}^k_{0:T}) = \\{1, \text{for } \kappa=\arg\min\_j \lVert z^k-e\_j\lVert\_2; 0\text{ otherwise.}\\}$ is a categorical distribution and becomes deterministic after the quantization operation. By adding noise, we reintroduce variability that allows the model to better varying data, thus improving the model's ability to generalize and increase robustness in different scenarios.

Reviewer comment (Evaluation):
Are there any significant differences between the two?

Author response:
**Q2: Is there any significant difference between L and state in IMDP? Can L be regarded as an unobserved state.**

**A2**: Mathematically, our $L$ is a latent confounder of all states. All existing work for MDP with latent variables have not considered this specfical case. Existing results cannot be directly applied. The details are as follows:

- The state $ S_t $ usually refers to the specific situation of environment where the agent is present at timestep $ t $, and it only directly influences the subsequent state $ S_{t+1} $.

- In contrast, $ L $  directly influences different states across all timesteps throughout the transition process.  $ L $ remains constant over time for each individual, while it may vary between individuals.  Mathematically, it is a latent confounder of all states.

To the best our knowledge, we failed to find any works discuss such a latent confounder in the Markov decision process. Therefore existing work for MDP with latent variables cannot be directly applied here. These methods typically require that latent variables change over time [1-4]. The identifiability of having a latent confounder L of all states has not been explored in previous research.



[1] Zhang, Xuezhou, et al. "Efficient reinforcement learning in block MDPs: A model-free representation learning approach." International Conference on Machine Learning. PMLR, 2022.

[2] Guo, Zhaohan Daniel, et al. "Bootstrap latent-predictive representations for multitask reinforcement learning." International Conference on Machine Learning. PMLR, 2020.

[3] Delgrange, Florent, Ann Nowe, and Guillermo A. Pérez. "Wasserstein Auto-encoded MDPs: Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees." arXiv preprint arXiv:2303.12558 (2023).]

[4] Zhang A, Lyle C, Sodhani S, et al. Invariant causal prediction for block mdps[C]//International Conference on Machine Learning. PMLR, 2020: 11214-11224.

Reviewer comment (Theory):
Can L be regarded as an unobserved state?

Author response:
**Q2: Is there any significant difference between L and state in IMDP? Can L be regarded as an unobserved state.**

**A2**: Mathematically, our $L$ is a latent confounder of all states. All existing work for MDP with latent variables have not considered this specfical case. Existing results cannot be directly applied. The details are as follows:

- The state $ S_t $ usually refers to the specific situation of environment where the agent is present at timestep $ t $, and it only directly influences the subsequent state $ S_{t+1} $.

- In contrast, $ L $  directly influences different states across all timesteps throughout the transition process.  $ L $ remains constant over time for each individual, while it may vary between individuals.  Mathematically, it is a latent confounder of all states.

To the best our knowledge, we failed to find any works discuss such a latent confounder in the Markov decision process. Therefore existing work for MDP with latent variables cannot be directly applied here. These methods typically require that latent variables change over time [1-4]. The identifiability of having a latent confounder L of all states has not been explored in previous research.



[1] Zhang, Xuezhou, et al. "Efficient reinforcement learning in block MDPs: A model-free representation learning approach." International Conference on Machine Learning. PMLR, 2022.

[2] Guo, Zhaohan Daniel, et al. "Bootstrap latent-predictive representations for multitask reinforcement learning." International Conference on Machine Learning. PMLR, 2020.

[3] Delgrange, Florent, Ann Nowe, and Guillermo A. Pérez. "Wasserstein Auto-encoded MDPs: Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees." arXiv preprint arXiv:2303.12558 (2023).]

[4] Zhang A, Lyle C, Sodhani S, et al. Invariant causal prediction for block mdps[C]//International Conference on Machine Learning. PMLR, 2020: 11214-11224.

Reviewer comment (Theory):
Lack warranty in some major claims. Though without references, I can understand some reasons described by the authors why only considering discrete representations, since some underlying latence of humans could be discrete such as gender, and continuous representations can be translated into discrete. However, the authors state that the discrete formulation allows for a more nuanced representation of individualized latent, which is lack theoretical and empirical evidence. If that is an assumption, I would also recommend such strong assumption to be clearly formulated and described in main content, given it plays a key role in the methodology.

Author response:
**Q1: Lack warranty on why discrete formulation is a more nuanced representation.** 

**A1:** Thank you for pointing that out. We are sorry for the confusion and we have removed this part. A new revised version is provided in Section 4.3:

> Group determinacy indicates that for a given task, the individual-specific factors are finite. This is a realistic assumption in numerous practical situations. For example, in hypertension diagnosis, patient groups may be formed based on factors such as age, gender, or medical history. Each group would demonstrate unique health outcomes and patterns, driven by their respective latent individual-specific factors $l_k$.

> we can divide the continuous latent confounder $L$ into small segments and represent each with a discrete value $l_k$.

Reviewer comment (Theory):
Design of methodology is not well motivated, and some parts lack details. For example, (a) The motivation of choosing VAE to estimate individual-specific factors, and the motivation of state reconstruction are not clear. The part of feature extraction is not introduced in overview, and the definition of features, and the reason of utilizing feature extraction, are not clear. (b) The proposed method relies on subgrouping of individuals to obtain embedding dictionary, while the part of subgrouping is highly vague in paper and lack thorough motivation and details. For example, (a). There is no reference regarding the motivation of such design, and reasons of choosing the specific subgrouping measure. (b) q is not included in pseudocode. Specific measures to obtain q and how q is determined in experiments are not clear.

Author response:
**Q2: The motivation of choosing VAE to estimate individual-specific factors is not clear.**  

**A2:** Thank you. Our proposed framework is carefully customized to meet the requirements of the identifiability theorem. According to the definition of identifiability (detailed in Appendix A.3.1), if a learned generative model matches the true model in terms of joint distribution, it would help establish the identifiability of latent individual-specific factors.

This motivates our work to use a generative model and to estimate confounders by fitting the learned distribution to the true observed distribution. Specifically, we use an encoder to infer the low-dimensional latent confounder, and a decoder to mimic the data generation process. This naturally leads us to use the VAE framework. We also found that the performance of the traditional VAE is not sufficient for recovering the latent individual-specific factor, as shown in the ablation study in Section 5.1. Thus, we have developed an individual-specific factor estimation framework to recover the latent factors.

**Q3: The motivation of state reconstruction is not clear.**  

**A3:**  According to the identifiability definition, the better the data generation process is learned, the better $\hat{l}_k$ would be recovered. To reconstruct individualized transition processes, we design a conditional decoder to reconstruct the state. The discrepancy between the reconstructed state and the original state provides the reconstruction loss for the optimization of the entire framework.

**Q4: The motivation of feature extraction.**  

**A4:**  Sorry for the confusion. "Feature" refers to the output of the encoder before discretization. It can be viewed as a continuous representation of the latent individual-specific factors. To avoid any misunderstanding, we have replaced the original term "feature" with "latent low-dimensional continuous representation."

**Q5: Subgrouping is highly vague in paper and lacks thorough motivation and details.**  

**A5:** Regarding the subgroup you mentioned, we guess you are referring to the embedding dictionary in the vector quantization layer.

The concept of the vector quantization layer originates from the paper [1]. his process converts continuous representations into a discrete latent space by identifying the closest vector in a predefined embedding dictionary $E$. The quantized vector $q$, denoted as $\hat{l}$, represents the embedding $e_i$ that is closest to the encoder output $z$, mathematically expressed as $\hat{l}\_k=\arg\min\_{e\_i}\lVert z^k-e\_i\lVert\_2$. This measurement is commonly employed in other related works focused on quantization [2-3].

In the experiment, $\hat{l}_k$ is determined through the process of vector quantization, and the learning process is optimized based on the criterion of minimizing the Euclidean distance between $z$ and $e_i$. We have revised Section 4 and the pseudocode to make this clearer and to avoid any confusion.

> Considering the group determinacy assumption, the estimated latent individual-specific factor should be discrete. We then embed a quantization layer to discretize the continuous latent representation $z^k$ to the latent individual-specific factor $\hat{l}_k$.

> The assignment of a dictionary vector $e_i$ to $z^k$ can be realized by finding the nearest neighbor in the dictionary: $\hat{l}\_k=\arg\min\_{e_i}\lVert z^k-e_i\lVert\_2$, where $\hat{l}\_k$ represents the quantized vector that is the closest embedding $e_i$ to the continuous representation $z^k$.

[1] Van Den Oord A, Vinyals O. Neural discrete representation learning[J]. Advances in neural information processing systems, 2017, 30.

[2] Yu J, Li X, Koh J Y, et al. Vector-quantized image modeling with improved vqgan[J]. arXiv preprint arXiv:2110.04627, 2021.

[3] Gupta A, Mukhopadhyay R, Balachandra S, et al. Towards Generating Ultra-High Resolution Talking-Face Videos with Lip synchronization[C]//Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. 2023: 5209-5218.

Reviewer comment (Experiments):
Design and details of experimental are not clear, which hinders the evaluation of proposed method and reproducibility. Some important details, such as: for synthetic experiments, the goal of the task, how the synthetic datasets are conducted, definitions of states, actions, and rewards, why only 100 samples are generated, how the unobserved latent L and noise are simulated, as well as the reasons of design such synthetic environment without considering existing broadly used environments (e.g., Adroit, Mujoco), especially human-involved ones. For persuasion for good environment, objectives, states, actions, and rewards are not formally defined. Definition of relevant information from the dialogues using the BERT is not provided.

Author response:
**Q6: Design and details of the experiment are not clear in the synthetic experiment.** 
 
**A6:** Thank you for pointing this out. We have updated Section 5 with more experimental details for clarity. Here is what we have added: the task's goal, the method of creating synthetic datasets, the definition of states and actions, and a concise summary and emphasis on how we simulated the unobserved latent factor $L$ and noise in Section 5.1. We did not introduce a reward because it is not essential for evaluating the latent individual-specific factor. Figure 3 is corrected and the placeholder is removed. For additional details, please refer to Appendix D. Specifically, the modifications are as follows:

> We manually created time series data based on the post-nonlinear model and customize different types of latents with the required assumptions satisfied or violated. The data generation process is modeled as $\vec{s}\_t=f\_2(f\_1(\vec{s}\_{t-1},\vec{a}\_{t-1},L),\epsilon\_t)$, where $f\_1$ represents the nonlinear effect, and $f\_2$ denotes the invertible post-nonlinear distortion on $\vec{s}\_t$. The state $\vec{s}\_t\in\mathbb{R}^3$ and actions $\vec{a}\_t\in\mathbb{R}^2$ are initially generated randomly following a uniform distribution, and the noise is modeled as a Gaussian distribution. Three types of $L$ are generated to satisfied or violate our assumptions. Case 1: $L$ is a discrete latent scalar follows the categorical distribution $\text{Cat}(0.1,0.2,0.3,0.4)$. Case 2: $L$ are multiple discrete latent variables follow the categorical distributions $\text{Cat}(0.1,0.9)$, $\text{Cat}(0.2,0.3,0.5)$, $\text{Cat}(0.1,0.2,0.3,0.4)$ respectively. Case 3: $L$ are multiple continuous latent variables follow the Gaussian distribution $\mathcal{N}(0,1)$, uniform distribution ${\rm Uniform}(0,1)$, and exponential distribution ${\rm Exp}(1)$.

**Q7: The reasons why design such synthetic environments without considering existing broadly used environments, especially human-involved ones.**

**A7:** To empirically validate the identifiability results for the individual-specific factor, we need a dataset with a known ground truth that is guaranteed to satisfy the proposed assumptions, including the presence of individual-specific factors. Unfortunately, this condition cannot be directly met by the existing widely-used environments. Therefore, this serves as one of the motivations for using a synthetic dataset for our validation purposes.

Additionally, we can easily modify the synthetic environments to explore what happens when specific assumptions are violated and how it affects the identifiability performance, providing more flexibility to either violate or satisfy our assumptions. Specifically, in Section 5.1, we customized three different types of latent variables involving both continuous/discrete and single/multiple confounders. If we use existing environments, we cannot customize and generate the desired dataset.

It is worth mentioning that we have also modified the Pendulum environment, which is a widely-used RL benchmark, to satisfy the individualized property and generate a population dataset.

**Q8: Why only 100 samples are generated in the synthetic environment.**

**A8:** The key to our sample sufficiency assumption is that the trajectory length should be greater than the number of groups; the number of individuals is less important. In the experiment, we set the number of individuals as 100 to demonstrate the effectiveness of the proposed method. Moreover, to address your concern, we will vary the number of samples to verify the effectiveness of our method. The experiments are still running due to the limited time, and we will update the results in openreview as soon as possible.

**Q9: For persuasion for good environment, objectives, states, actions, and rewards are not formally defined.**

**A9:** Thank you. In the revised version, we have defined the objectives, states, actions, and rewards in the persuasion for good dataset in Section 5.1. 

> We further evaluate the effectiveness of the estimation framework on a real-world, open-source dataset, the Persuasion For Good Corpus~\citep{mydataset,wangetal2019persuasion}. It is a collection of dialogues focused on charitable donations, and each dialog involves a persuader trying to convince a persuadee to donate to a charity, where all participants have undergone personality assessments and have corresponding individualized personalities. The state is the persuadee's response, the action is the persuader's utterance, and the reward is the final donation in each dialog. The goal is to make an individualized policy for each persuader to successfully persuade the persuadee to make a charitable donation. To achieve this, we first need to identify the latent personality of each individual from the offline dataset.

Reviewer comment (Presentation):
Caption of Figure 3(a): “XX” looks like a placeholder.

Author response:
**Q6: Design and details of the experiment are not clear in the synthetic experiment.** 
 
**A6:** Thank you for pointing this out. We have updated Section 5 with more experimental details for clarity. Here is what we have added: the task's goal, the method of creating synthetic datasets, the definition of states and actions, and a concise summary and emphasis on how we simulated the unobserved latent factor $L$ and noise in Section 5.1. We did not introduce a reward because it is not essential for evaluating the latent individual-specific factor. Figure 3 is corrected and the placeholder is removed. For additional details, please refer to Appendix D. Specifically, the modifications are as follows:

> We manually created time series data based on the post-nonlinear model and customize different types of latents with the required assumptions satisfied or violated. The data generation process is modeled as $\vec{s}\_t=f\_2(f\_1(\vec{s}\_{t-1},\vec{a}\_{t-1},L),\epsilon\_t)$, where $f\_1$ represents the nonlinear effect, and $f\_2$ denotes the invertible post-nonlinear distortion on $\vec{s}\_t$. The state $\vec{s}\_t\in\mathbb{R}^3$ and actions $\vec{a}\_t\in\mathbb{R}^2$ are initially generated randomly following a uniform distribution, and the noise is modeled as a Gaussian distribution. Three types of $L$ are generated to satisfied or violate our assumptions. Case 1: $L$ is a discrete latent scalar follows the categorical distribution $\text{Cat}(0.1,0.2,0.3,0.4)$. Case 2: $L$ are multiple discrete latent variables follow the categorical distributions $\text{Cat}(0.1,0.9)$, $\text{Cat}(0.2,0.3,0.5)$, $\text{Cat}(0.1,0.2,0.3,0.4)$ respectively. Case 3: $L$ are multiple continuous latent variables follow the Gaussian distribution $\mathcal{N}(0,1)$, uniform distribution ${\rm Uniform}(0,1)$, and exponential distribution ${\rm Exp}(1)$.

Reviewer comment (Presentation):
The font size in the plots is too small to read.

Author response:
**Q2: The font size in the plots is too small to read.**

**A2:** Thank you for pointing this out. We have increased the font size and clarity of all figures to improve readability in the revised manuscript, and Figure 2 has been redrawn to better illustrate the framework's structure.

Reviewer comment (Experiments):
I'm not sure that pendulum control counts as a real-world application.

Author response:
**Q3: I'm not sure that pendulum control counts as a real-world application.**

**A3:** Thank you for pointing this out. We have changed the section name to 'Evaluation on Policy Learning Framework' to avoid misunderstanding.

Reviewer comment (Novelty):
The contribution of the paper is quite small. Much of the contribution is simply incorporating a VAE into an RL system.

Author response:
**Q4: The contribution of the paper is quite small. Much of the contribution is simply incorporating a VAE into an RL system.**

**A4:** Thanks. We would like to emphasize that our contributions go far beyond the integration of a VAE into a RL system. 

To the best of our knowledge, none of the existing works consider a time-invariant latent confounder that influences each state in the transition process. We are the first to model such processes and provide an identifiability result.

Specifically, existing works [1-3] usually focus on latent variables that are time-varying. When they are time-varying, they can benefit from many recent advances in nonlinear ICA to achieve strong identifiability results [4-5]. However, the identifiability of time-invariant latent confounders, though not well-studied, has numerous applications. For example, in healthcare, time-invariant latent confounders could be gene variations, birth conditions, or genetic predispositions, which influence the state transition process of health status.


In this work, we aim to bridge this gap. We seek to answer the important question of which types of transition processes can have a theoretical identifiability guarantee. We demonstrate that when the individual-specific factors are finite, our method ensures the identifiability of the entire latent state-transition process, even when the transition processes are nonlinear. This establishes novel theoretical insights for learning state-transition processes with latent individual-specific factors.

Furthermore, motivated by the identifiability results, we have designed a practical generative framework, aligning with our theoretical assumptions, **which is very different from VAE**. Specifically,

1) To satisfy the group determinacy assumption, which assumes there is a finite number of L, we utilize a quantized encoder where the embedding dictionary contains only a finite number of L.

2) To satisfy the sequential dependency assumption, which assumes the generative process where $s_t$ is determined by $s_{t-1}$, $a_{t-1}$, L and noise, we model the noise with an MLP layer. Then, a conditional encoder, along with a noise layer, is designed to model the data generation process.

To enable readers to quickly grasp the essence of our contribution, we have revised the Introduction as follows:

> We propose the Individualized Markov Decision Processes (iMDPs) framework, a novel approach that integrates individual-specific factors into decision-making processes. We model these factors as latent confounders. We allow them to influence each state in the decision process and to vary across different individuals. This framework has many real-world applications.

> Our framework has theoretical guarantees. We show that when individual-specific factors are finite, our method ensures the identifiability of whole latent state-transition processes, even when the transition processes are nonlinear. This establishes novel theoretical insights for learning state-transition processes with latent individual-specific factors. Additionally, for scenarios with non-finite individual-specific factors, we show that categorizing them into finite groups has little impact on the empirical results.

To the best of our knowledge, this is the first work to provide a theoretical guarantee for the identification of latent individual-specific factors from observed data.

> We propose a practical generative-based method that can effectively infer latent individual-specific factors. Empirical results on both synthetic and real-world datasets demonstrate the method's effectiveness not only in inferring these factors but also in learning individualized policies.

[1] Zhang A, Lyle C, Sodhani S, et al. Invariant causal prediction for block mdps[C]//International Conference on Machine Learning. PMLR, 2020: 11214-11224.

[2] Delgrange, Florent, Ann Nowe, and Guillermo A. Pérez. "Wasserstein Auto-encoded MDPs: Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees." arXiv preprint arXiv:2303.12558 (2023).

[3] Guo, Zhaohan Daniel, et al. "Bootstrap latent-predictive representations for multitask reinforcement learning." International Conference on Machine Learning. PMLR, 2020.

[4] Hyvarinen A, Morioka H. Unsupervised feature extraction by time-contrastive learning and nonlinear ica[J]. Advances in neural information processing systems, 2016, 29.

[5] Yao W, Sun Y, Ho A, et al. Learning temporally causal latent processes from general temporal data[J]. arXiv preprint arXiv:2110.05428, 2021.

Reviewer comment (Novelty):
The novelty is somewhat overclaimed. There are many works, such as [1-5], on learning latent/factored MDPs, which are very similar to the so-called individualized MDPs. However, the paper does not clearly clarify the difference and contribution.

Author response:
**Q1: As there are many similar works [1-5], the novelty is somewhat overclaimed.** 

**A1:** We apologize for the confusion. Our method is fundamentally distinct from existing approaches [1-5], as it is based on a different latent transition process. Both our theoretical results and practical methods are grounded in this unique aspect of the latent transition process, making our approach different from existing methods. 

**Here are some key differences in the latent transition process.**

- In our work, we consider a latent confounder (individual-specific factor) that influences each state in the state transition process, and the proposed individualized transition process is motivated by numerous applications. We provide theoretical results that when individual-specific factors are finite, our method ensures the identifiability of the entire latent state-transition process, even in cases of nonlinear transitions. This establishes novel theoretical insights for learning state-transition processes with latent individual-specific factors.

- However, Factored MDP [1] and Factored Non-stationary MDP [3] assume there are **no unobserved confounders** in the state transitions. Block MDP [2], POMDP [4], and Latent MDP [5] consider **latent states/spaces, but such latent factors do not influence each state in the transition process**. The aforementioned settings differ significantly from our work.

**Here we would like to provide more details regarding on the difference.** 

**To the best of our knowledge, none of the existing works consider a time-invariant latent confounder that influences each state in the transition process. We are the first to model such processes and provide an identifiability result.**

Specifically, existing works [1-3] usually focus on latent variables that are time-varying. When they are time-varying, they can benefit from many recent advances in nonlinear ICA to achieve strong identifiability results [4-5]. However, the identifiability of time-invariant latent confounders, though not well-studied, has numerous applications.

(For example, in healthcare, two patients diagnosed with hypertension may have different health outcomes (state $s_t$) after receiving the same treatment (action $a_t$), influenced by the time-invariant latent confounders ($L$) such as gene variations, birth conditions, or genetic predispositions. If we could observe these predispositions, it would greatly assist physicians in predicting how patients will respond to specific treatments and facilitate personalized healthcare. However, in most cases, these individual-specific factors remain latent and are difficult to diagnose.)

The theoretical identification of such latent confounders and the development of an estimation method to extract them from observational data is of immense practical value. Understanding such a transition process accelerates the development of individualized machine learning.

In this work, we provide a theoretical understanding of such a transition process.We demonstrate that when individual-specific factors are finite, our method ensures the identifiability of the latent individual-specific factors, even when the transition processes are nonlinear. This establishes novel theoretical insights for learning state-transition processes with latent individual-specific factors.

**Furthermore, motivated by the identifiability results, we have designed a practical generative framework, aligning with our theoretical assumptions, which is very different from VAE**. Specifically,
- To satisfy the group determinacy assumption, which assumes there is a finite number of L, we utilize a quantized encoder where the embedding dictionary contains only a finite number of L.
-  To satisfy the sequential dependency assumption, which assumes the generative process where $s_t$ is determined by $s_{t-1}$, $a_{t-1}$, L and noise, we model the noise with an MLP layer. Then, a conditional encoder, along with a noise layer, is designed to model the data generation process.

[1] Zhang A, Lyle C, Sodhani S, et al. Invariant causal prediction for block mdps[C]//International Conference on Machine Learning. PMLR, 2020: 11214-11224.

[2] Delgrange, Florent, Ann Nowe, and Guillermo A. Pérez. "Wasserstein Auto-encoded MDPs: Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees." arXiv preprint arXiv:2303.12558 (2023).

[3] Guo, Zhaohan Daniel, et al. "Bootstrap latent-predictive representations for multitask reinforcement learning." International Conference on Machine Learning. PMLR, 2020.

[4] Hyvarinen A, Morioka H. Unsupervised feature extraction by time-contrastive learning and nonlinear ica[J]. Advances in neural information processing systems, 2016, 29.

[5] Yao W, Sun Y, Ho A, et al. Learning temporally causal latent processes from general temporal data[J]. arXiv preprint arXiv:2110.05428, 2021.

Reviewer comment (Novelty):
This paper seems to claim that there is no prior art in this field; however, the reviewer believes that MARL (Multi-Agent Reinforcement Learning) methods and latent MDP (Markov Decision Process) methods could also be applied to the setting described in this paper.

Author response:
**Q2: Multi-Agent Reinforcement Learning methods and latent Markov Decision Process methods could also be applied to the setting described in this paper.**

**A2:** It is hard to apply Multi-Agent Reinforcement Learning (MARL) and Latent Markov Decision Processes (MDPs) in our cases. The key points are summarized as follows:

Regarding to Multi-Agent Reinforcement Learning

**MARL typically involves multiple agents interacting in a shared environment, but in our work, each individual interacts with a unique environment. Therefore, MARL methods cannot be directly applied to our setting.** Specifically, in MARL, the actions of one agent can influence the outcomes of others, and agents jointly contribute to the dynamics of the environment. In contrast, in our work, each individual has a distinct environment, and there is no direct interaction between different individuals. The state transition processes are implicitly influenced by a latent individual-specific factor L.

**If we tweak the MARL setting to allow different agents to have individualized environments, which could potentially be adapted to our setting, the agents would be trained independently.** In that case, the training would require more data, as all agents are trained separately in each environment, making it impossible to generalize across different environments in general. For example, the policy would need to be retrained when a new agent in a new individualized environment is introduced.

Regarding to Latent Markov Decision Processes



On the other hand, **in Latent MDP, latent variables (states) change with respect to time. However, in our work, we have a latent confounder L that influences each state in the transition process. All existing work for MDP with latent variables have not considered this specific case. Existing results cannot be directly applied. The details are as follows:

- The state $ S_t $ usually refers to the specific situation of environment where the agent is present at timestep $ t $, and it only directly influences the subsequent state $ S_{t+1} $.

- In contrast, $ L $  directly influences different states across all timesteps throughout the transition process.  $ L $ remains constant over time for each individual, while it may vary between individuals.  Mathematically, it is a latent confounder of all states.

To the best our knowledge, we failed to find any works discuss such a latent confounder in the Markov decision process. Therefore existing work for MDP with latent variables cannot be directly applied here. These methods typically require that latent variables change over time [1-4]. The identifiability of having a latent confounder L of all states has not been explored in previous research.


[1] Zhang, Xuezhou, et al. "Efficient reinforcement learning in block MDPs: A model-free representation learning approach." International Conference on Machine Learning. PMLR, 2022.

[2] Guo, Zhaohan Daniel, et al. "Bootstrap latent-predictive representations for multitask reinforcement learning." International Conference on Machine Learning. PMLR, 2020.

[3] Delgrange, Florent, Ann Nowe, and Guillermo A. Pérez. "Wasserstein Auto-encoded MDPs: Formal Verification of Efficiently Distilled RL Policies with Many-sided Guarantees." arXiv preprint arXiv:2303.12558 (2023).]

[4] Zhang A, Lyle C, Sodhani S, et al. Invariant causal prediction for block mdps[C]//International Conference on Machine Learning. PMLR, 2020: 11214-11224.

Reviewer comment (Theory):
The theoretical result requires unrealistic assumptions. The group determinacy, sample sufficiency, and sequential dependency assumptions are too strong to be realistic. Thus, the obtained theoretical results, though seemingly novel, have no significant implications. The reviewer cannot derive insights from this theorem.

Author response:
**Q4: The theoretical result requires unrealistic assumptions.** 

**A4:** We believe this is a misunderstanding. Our assumptions sometimes may not be satisfied but still realistic. 

To avoid confusion, we have added a detailed explanation of our assumptions in Appendix A.2. Specifically:

> **Group determinacy** indicates that for a given task, the individual-specific factors are finite. This is a realistic assumption in numerous practical situations. For instance, in healthcare, consider two patients diagnosed with hypertension. Even if they receive the same treatment (or action $a_t$), they might experience different health outcomes (or state $s_t$). These differences can be attributed to latent factors like their medical history (or latent individual-specific factors $L$), which are essentially finite records of past medical events and treatments. In this case, we can group patients based on their history, creating categories such as those with no history of hypertension and those with a history of hypertension. Each group would demonstrate unique health outcomes and patterns, driven by their respective individual-specific factors $l_k$. 


> **Sequential dependency** is commonly seen in other related works, such as Temporal Causal Mechanisms in Block MDP [1]. It ensures that only previous information (such as state and action) influences the current state. In the case of individualized healthcare, this assumption implies that a patient's latent factors (such as medical history), current health state and current treatment decisions determine their next health state. 

> **Sample Sufficiency** indicates that sufficient data is needed to achieve identifiability, which should be a common assumption. For example, in nonlinear ICA using auxiliary variables [2], it assumes that there should be at least 2n+1 values for auxiliary variables to encourage variability; for disentanglement with minimal change [3], it assumes that there should be at least 2n+1 domain embeddings to ensure identifiability. Intuitively, without having data to provide us with related information about parameters, it is impossible to “determine” the value of these parameters.

Note that even though the required assumptions may not always hold in practice, our method can still encourage identifiability in different cases, namely: (1) multiple independent latent confounders, (2) multiple independent continuous latent confounders, and (3) insufficient samples. A detailed discussion has been added in Section 4.3, and these propositions are empirically validated in the experiments in Section 5.1. Sepcifically,

> In the first case, we can treat multiple latents as one discrete latent with cardinality $\prod^m_{i=1}c_i$. Then, each combination of the original multiple latents can be uniquely represented as one level in the combined latent variable. In the second case, the empirical results in Section 5 show that our proposed framework can also encourage the identifiability of the continuous latents. This is because we can divide the continuous latent confounder $L$ into small segments and represent each with a discrete value $l_k$. This discretization approach improves computational simplicity and stability. In the third case, even when the sample sufficiency assumption is weakly violated, our framework still encourages the identification of the latent confounder, as verified in the experiment in Section 5.

[1] Zhang A, Lyle C, Sodhani S, et al. Invariant causal prediction for block mdps[C]//International Conference on Machine Learning. PMLR, 2020: 11214-11224.

[2] Hyvarinen A, Sasaki H, Turner R. Nonlinear ICA using auxiliary variables and generalized contrastive learning[C]//The 22nd International Conference on Artificial Intelligence and Statistics. PMLR, 2019: 859-868.

[3] Kong L, Xie S, Yao W, et al. Partial Identifiability for Domain Adaptation[J]. arXiv preprint arXiv:2306.06510, 2023.

Reviewer comment (Theory):
The symbol $L$ should depend on $k$ in Equation (1).

Author response:
**Q5. The symbol L should depend on k in Equation (1).**

**A5:** Thank you for pointing it out. We have updated Equation 1 and replaced $L$ with $l_k$.

Reviewer comment (Presentation):
To improve readability, background information about VAE should be provided before applying (VQ-)VAE in Section 4.

Author response:
**Q6. Background information about VAE should be provided before applying (VQ-)VAE in Section 4.**

**A6:** We have enriched Section 4 with additional background information on VAEs. To save space in the main paper, more detailed background on VAEs can be found in Appendix C. Specifically, the additions are cited as follows.

> The traditional VAE encodes the input data into a continuous latent space using probabilistic encoders and then reconstructs the input from this space using decoders, which is not aligned with our problem setting.

> Variational Autoencoders (VAEs) are a class of generative models in deep learning, adept at unsupervised learning of complex data distributions. Rooted in the framework of Bayesian inference, VAEs are designed to approximate probability density functions of input data. The architecture of a VAE consists of two primary components: an encoder $q_{\phi}(z|x)$ and a decoder $p_{\theta}(x|z)$. The encoder maps input data $x$ to a latent space, represented by a probability distribution, typically Gaussian, with parameters $\mu$ and $\sigma$ signifying the mean and standard deviation, respectively. The decoder reconstructs the input data from a sampled latent representation $z$.

Reviewer comment (Presentation):
The font size in Figure 3 is too small to be readable.

Author response:
**Q8. The font size in Figure 3 is too small to be readable.**

**A8:** Thank you for pointing that out. We have increased the font size in Figure 3 to make it easier to read.

Reviewer comment (Writing):
The difference between individualized MDPs, partially observed MDPs and multi-agent MDPs is unclear.

Author response:
**Q3: What is the difference between individualized MDPs, partially observed MDPs and multi-agent MDPs.**

**A3:**  We have made one-to-one comparisons as follows:

- **POMDP  vs.  Individualized MDPs** 
  - POMDPs are similar to Latent MDPs. However, they do not consider our case where there is a latent confounder $L$ that influences each state in the transition process.

  - **Incomplete Information vs. Latent Individual-Specific Factor.** POMDPs deal with situations where the agent has incomplete or uncertain information about the state of the environment. Such incomplete information is modeled as the hidden state space, which has a mapping to the observational space. Individualized MDPs, on the other hand, deal with fully observable states, but the observed state transition process is implicitly influenced by the latent individual-specific factor.
  - **Time-Varying Influence on the Observation vs. Time-Invariant Influence on the Observation** In POMDPs, the latent states are time-varying, while in Individualized MDPs, the latent individual-specific factor is time-invariant for each individual.
  -**Decision-making under Uncertainty vs. Decision-making for Individual Needs** POMDPs focus on scenarios where the agent must make decisions without complete information about the environment, while Individualized MDPs focus on adapting decisions to individual needs and preferences.

- **MAMDPs vs. Individualized MDPs**
  - **Agent Interaction vs. Individual Focus**: MAMDPs involve multiple agents whose actions affect each other, emphasizing the dynamics of all agents' interactions in one shared environment. In contrast, Individualized MDPs focus on optimizing decisions based on the unique characteristics of individual agents. The individuals share a common latent property but do not interact with each other.
  - **Collective Optimization vs. Individualized Optimization**: MAMDPs often require strategies that balance collective outcomes for all agents, while Individualized MDPs aim to optimize outcomes for each individual, considering their specific context or requirements.
(Please also refer **A2** for why method for MAMDPs hard to be applied to our cases).

Reviewer comment (Theory):
The optimal policy appears to be non-stationary in individualized MDPs, and it is unclear how this difficulty can be addressed.

Author response:
**Q10: The optimal policy appears to be non-stationary in individualized MDPs. How can this difficulty be addressed.**

**A10:** We guess the non-stationarity you mentioned refers to the policy varying across individuals, considering that changes in the latent factors lead to changes in the transition dynamics, whereas the individualized policy is stationary for each individual.

To address this issue, we have included the latent individual-specific factor as the policy input and adjusted the training objective accordingly to prioritize the optimization of individual policies. When a new individual is encountered, the agent simultaneously optimizes its policy and collects new data to estimate the latent factors. The details have been added in Section 4.2 and are as follows:

> After optimizing the individualized policy, the agent uses $\pi^*_k$ as a warm start, which is an initial starting point that helps speed up training, and continues training on the new individual. To achieve zero-shot transfer, we simultaneously optimize the individualized policy while interacting with the environment. Thus, adaptation involves a dual process: it fine-tunes the policy based on interactions with the new individual, and it actively collects data from these interactions. This data collection is critical for estimating the latent individual-specific factor of the new individual. During this phase, the policy is continuously adjusted and improved, making it better suited to each new individual.