Identifying Latent State-Transition Processes for Individualized Reinforcement Learning

Authors: biwei  huang
yu  yao
xinshuai  dong
songyao  jin
ucsd  boyang
sun  mbzuai
kazushi  ikeda
peter  spirtes
kun  zhang
mbzuai  cmu
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 21:29:38

Abstract: In recent years, the application of reinforcement learning (RL) involving interactions with individuals has seen significant growth. These interactions, influenced by individual-specific factors ranging from personal preferences to physiological differences, can causally affect state transitions, such as the health conditions in healthcare or learning progress in education. Consequently, different individuals may exhibit different state-transition processes. Understanding these individualized state-transition processes is crucial for optimizing individualized policies. In practice, however, identifying these state-transition processes is challenging, especially since individual-specific factors often remain latent. In this paper, we establish the identifiability of these latent factors and present a practical method that effectively learns these processes from observed state-action trajectories. Our experiments on various datasets show that our method can effectively identify the latent state-transition processes and help learn individualized RL policies.

Main paper: 
 introduction: 
 Reinforcement Learning (RL) [46] involves training agents to make decisions by interacting with the environment. The agent observes its current state, takes an action, and transitions to a new state with a reward. Such a sequence of moving from one state to another is known as a state-transition process.Individualized RL focuses on adapting the policy for each individual. It has recently seen increasing application in various sectors, including healthcare [89, 57, 21] , education [68, 2, 14] , and e-commerce [53, 87, 1] . The individual-specific factors [60] , which embed the unique characteristics of each individual, play a crucial role in causally influencing the transitions between states. Such factors range from individual preferences and past experiences to physiological differences. For example, in the realm of education, different learning styles can affect how two students with the same prior knowledge learn from a tutorial. In healthcare, differences in genetic makeup can affect how two hypertension patients respond to identical treatments. Understanding individual-specific factors is essential for designing better RL systems that provide more individualized and effective decisions [41, 28, 4, 65] . With knowledge of learning styles, the RL agents can recommend personalized tutorials, such as animated content for visual learners or hands-on exercises for kinesthetic learners. Similarly, in healthcare, such knowledge of genetic makeup can help agents suggest treatment plans tailored to their specific needs, leading to improved health outcomes. However, these individual-specific factors are not always observable, which poses challenges to understanding the Latent Individualized State-Transition (LIST) processes, as shown in Figure 1 (a).The latent individual-specific factors are unique (e.g., learning styles, genetic makeup, etc.) to each individual and have a time-invariant influence on the state-transition process. This raises the question: can we guarantee the identifiability of the latent factors?Such identifiability is easier to achieve if the observations are either i.i.d., or i.i.d. given side information (e.g., domain index, time index, etc.) by exploiting sparsity [97] , variability [47] , or functional complexity [43] . To the best of our knowledge, only a few studies have attempted to uncover the identifiability of latent factors from temporal observations. These methods focus on the time-varying latent factors, which is very different from our work on time-invariant latent factors. Specifically, existing works [86, 85, 7] assume the time-varying latent variables without considering the influence of actions in the generative process (see Figure 1(b) ). Factored MDPs [15] incorporate actions into the process but still assume that the latent factors change over time (see Figure 1(c) ). Thus the results from existing work cannot be applied in our setting. Intuitively, this is because the timeinvariant latent factors cannot provide the variability that many current methods rely on to achieve identifiability. It remains unknown how to derive the identifiability of the latent individual-specific factors, together with the latent state-transition processes, from the observed states and actions.Recent advances in finite mixture models [77, 67] have proven strong identifiability results by exploiting group information in nonparametric settings. By assuming that observations in the same group are known to come from the same component, the mixture of probability measures can be uniquely identified under proper assumptions. Inspired by these works, we establish the identifiability of the latent factors by leveraging group information from the data, making it easier to distinguish different underlying components. We propose both finite latent, nonparametric setting and infinite latent, parametric setting and develop a theoretically grounded framework that effectively learns these processes from observed state-action trajectories. Our contributions are summarized as follows:• We propose the Individualized Markov Decision Processes (iMDPs), a novel framework that integrates latent individual-specific factors κ into state-transition processes. We consider κ as a latent individual-specific factor, allowing it to influence each state in the decision process and to vary across different individuals. • Our work provides theoretical guarantees and new insights for learning state-transition processes with latent factors. When κ is finite, we consider two scenarios to derive the identifiability even if the processes are nonparametric. For infinite κ, we show that identifiability is guaranteed in the post-nonlinear case. 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 transitions. • We propose a practical generative-based method that can effectively estimate the latent individualspecific factors. Empirical results on various datasets demonstrate the method's effectiveness not only in inferring these factors but also in learning individualized policies. 
related work: 
 Individualized Machine-Learning Applications Recently, machine learning has created highly individualized solutions across various domains. In healthcare, algorithms support individualized interventions for physical activity, weight loss, and diabetes management [89, 57, 18, 17] . In finance, it provides accurate stock predictions for stock market activities [56] . Education is benefiting from individualized ICT systems that address the individual learning needs of students [16, 40] . Furthermore, transportation has seen the development of individualized car-following strategies [69] that improve driving safety and efficiency. Meanwhile, entertainment platforms such as YouTube and TikTok are using it to provide individualized video recommendations [6, 33] .Reinforcement Learning with Latent State-Transition Processes In the field of RL, various models explore the state transition dynamics with latent variables. One such approach is Partially Observable Markov Decision Processes (POMDPs) [64] , where the full information about the state is unknown. In POMDPs, observations are generated from the latent states, which do not match our individual latent setting. For example, block MDPs [92, 96] assume that there is a fixed and unknown mapping from observations to the latent states. Factored MDPs [35, 15] , which provide the partial identifiability of latent factors, assume that the latent factors evolve over time following a Markov process. On the other hand, there exists a piece of work focusing on estimating state transitions with time-invariant latent factors. Models such as contextual MDPs [29, 60, 65] , latent MDPs [51, 50] and multitask RL [73, 24] consider similar scenarios with our latent individual-specific factors. However, these works lack theoretical guarantees on the identifiability of the latent factors thus it is hard for them to guarantee individualized decision-making. 
problem formulation: 
 Consider a population with M individuals that can be divided into G groups, whereas the exact group membership is unknown. We introduce iMDP to model individualized decision-making, where observed individual uniqueness is represented by u, and latent group-level properties are embedded by individual-specific factors κ. Specifically, each individual has a unique value of u, and the cardinality of u equals M . In contrast, individuals in the same group share the same value of κ and have different values across different groups, and the cardinality of κ equals G. For each individual, the value of κ is determined and has a time-invariant influence on the state-transition process. Suppose all individuals share the same state and action spaces. The iMDP is defined as follows.Definition 3.1 (iMDP). An iMDP consists of a tuple ⟨S, A, R, {s m 0 } M m=1 , {u m } M m=1 , {T m } M m=1 ⟩, where M is the number of individuals; S and A are the state and action spaces, respectively; R ∈ R is the immediate reward received after transitioning from s to s ′ via a, i.e., r = R a (s, s ′ ) for current state s ∈ S, new state s ′ ∈ S and action a ∈ A.For the m th MDP, u m is the unique index identifying each individual. s m 0 is the individualized initial state. T m ∈ R |S|×|A|×|S| is the individualized state transition probability, i.e., T m := P m (s ′ |s, a, κ m ). Here κ m is the latent individual-specific factor with cardinality G. Thus, the joint distribution of any adjacent state-action pairs (s, a, s ′ ) can be specified by u and κ as: P(s, a, s ′ |u) = P(s ′ |s, a, κ)P(s, a|u).(EQUATIONData Generation Process Here we introduce the latent individualized state-transition processes based on iMDP. For individual m, the observed states s m t satisfy the following generation process:EQUATIONwhere s m t = (s m 0,t , . . . , s m ds,t ) ⊤ ∈ R ds represents the state, and a m t = (a m 0,t , . . . , a m da,t ) ⊤ ∈ R da the action at time t, with d s and d a as the dimensions of state and action, respectively. ϵ m i,t denotes independent noise term. κ characterizes the group-level property across individuals, and the transition function f is identical across individuals, which is consistent with Eq. ( 1). During interaction with the environment, the trajectory τ m = {s m 0 , a m 0 , s m 1 , . . . , s m T } is recorded as a sequence of observed state-action tuples, where T denotes the trajectory length.Objectives In this work, we focus on the individualized RL agents with latent state-transition processes. Our objectives are twofold: 1) to identify latent individual-specific factors κ from observed trajectories, and 2) to derive individualized policies for each agent and realize policy adaptation for newcomers. Consider the example of hypertension diagnosis in healthcare. Treating all patients identically may lead to varied outcomes due to the dynamics of state transitions influenced by latent κ. Therefore, accurate identification of κ from the population provides crucial dynamic background knowledge. Once κ is uncovered, we can categorize patients into different groups and provide individualized treatments for each patient, which is consistent with our second goal.Consider a factor analysis model with p observations for each of n individuals and k common factors (k < p). The relationship is defined by the factor loading matrix L ∈ R p×k and the factor matrix F ∈ R k×n . Specifically,EQUATION)where X ∈ R p×n is the observation matrix and ε ∈ R p×n is the error term matrix. The factor loading matrix L and the factor matrix F are unique up to an orthogonal transformation. Specifically, for any orthogonal matrix Q, if we set L ′ = LQ and F ′ = Q T F , the transformed matrices L ′ and F ′ also satisfy the model criteria. 
identifiability analysis: 
 We consider two conditions that ensure the identifiability of latent individualized state-transition processes using either (1) group determinacy assumption or (2) functional constraints. The corresponding identifiability is established in the following theorems.Finite Latent Condition Suppose the value of κ is finite; we first provide the definition of groupwise identifiability. For the detailed assumptions discussion and proofs, please see Appendix B. Definition 4.1 (Group-wise Identifiability). Let {τ m } M m=1 be sequences of observed states and actions collected from G groups under a fixed policy, following the true latent individualized statetransition processes described in Eq. (2) . A learned generative model ( f , κ, ε) is observational equivalent to (f, κ, ϵ) if the joint distribution P f ,κ,ε (s, a, s ′ ) matches P f,κ,ϵ (s, a, s ′ ) everywhere. We say that the latent individualized state-transition processes are group-wise identifiable if observational equivalence can always lead to the identifiability of latent individual-specific factors across the population up to the invertible transformation g:EQUATIONAssumption 4.1 (Group Determinacy). Consider a finite mixture model G g=1 π g δ κg (κ), where π g represents mixing proportions with G g=1 π g = 1, and δ κg is the Dirac function centered at κ g . Each unique value of κ corresponds to a specific group in the population, with δ κg (κ) = 1 if κ = κ g and 0 otherwise, and the number of individuals per group is greater than 2G -1.Assumption 4.1 indicates that identifiability can be derived from the finite mixture model perspective using group information. We consider two scenarios to establish identifiability under finite latent conditions. Theorem 4.1 considers finite samples and specifies a minimum trajectory length with assumptions on initial states {s m 0 } M m=1 . Theorem 4.2 guarantees asymptotic identifiability for sufficiently long trajectories without initial state constraints. Theorem 4.1. Assume the LIST processes in Eq. (2) . Suppose the distributions of initial states within the same groups are the same for all individuals. Under Assumption 4.1, the identifiability of the individual-specific factor κ is guaranteed. Theorem 4.2. Assume the LIST processes in Eq. (2) . Suppose the distribution of initial state varies across individuals and the trajectory length is sufficiently long, i.e., there exist two different individuals in the same group have overlap condition P(s, a|u = u i ) = P(s, a|u = u j ), i ̸ = j. Under Assumption 4.1, the identifiability of κ is asymptotically guaranteed.The following theorem shows that under certain functional constraints, the identifiability of individual-specific latent variables can be extended to multiple and even infinite latent factors. Specifically, we consider the post-nonlinear temporal model [93] and allow multiple instances of κ to influence the state transition dynamics. The identifiability and the cardinality is decided upon the rank conditions of specific covariance submatrices derived from the observed data. In addition, the empirical results in Section 6 show that even when there are multiple latent factors with infinite cardinality, our estimation framework (see Section 5) still encourages the identification of latent factors. For the detailed model description and proof, please see Appendix C. Theorem 4.3. Consider a trajectory collected from the post-nonlinear temporal model (Definition C.1) with d s -dimensional observed states over time t = 1, . . . , T . Let m latent group factors κ j , j = 1, . . . , m, have direct causal influence on all states, and S t = {s 1,t , s 2,t , . . . , s n,t } represent the set of all state variables at time t. These latent factors, as well as the state-transition process, can be identified if and only if for every i = m + 2, . . . , T -(m + 1), there exist pairs of minimal rank sets (Definition C.2) (A i , B i ), defined as A i = R i,i -and B i = R i,i + , where i -< i < i + , that satisfy:• (Rank Deficiency for Identification) In addition to S i shared by A i and B i , each subset should include a randomly selected set of m + 1 additional state variables. If the covariance matrices Σ A i ,B i exhibit a consistent rank r (where r > d s ) across all distinct indices i, this consistency implies the existence of latent group factors within the system.• (Quantification of Latent Factors) Once identification is established, the count of latent group factors m can be inferred from the rank deficiency of Σ A i ,B i relative to the dimensionality of the observed variables, specifically given by m = rank(Σ A i ,B i ) -d s . 
estimation and policy learning framework: 
 We propose a two-stage approach to generate individualized policies. Our method achieves two objectives: (1) constructing an estimation framework to recover the latent group factors from a collection of individual trajectories, and (2) implementing individualized policy learning to facilitate policy adaptation for new individuals.Overview The proposed method is carefully designed to meet the requirements of the identifiability theorems. As specified in Definition 4.1, identifiability is ensured iff observational equivalence can always lead to latent factor equivalence. This motivates us to use a generative model to achieve latent factor estimation and ensure the learned distribution closely aligns with the true observed distribution.Theorems 4.1 and 4.2 provide the guarantee that such an alignment is asymptotically accurate. We adopt the variational autoencoder [45] architecture to estimate latent group factors and classify individuals into different groups in an unsupervised manner. The proposed estimation framework supports the identifiability theorem discussed in Section 4. As demonstrated in Figure 2 , the sequence of individual trajectories with the required size is represented in the discrete embedding space, which is consistent with the assumptions proposed in the theorems. A detailed pseudocode is provided in Appendix I, and a comprehensive realization of each component is available in Appendix H.Latent Inference via Quantized Encoding The group determinacy assumption suggests the existence of the latent individual-specific factor κ. Since κ is time-invariant and influences each state in the transition process, we first use a sequential encoder to capture the high-level representation z m , based on the input from all states across each trajectory s m 0:T . We then apply a vector quantization layer [76] to discretize the latent space and estimate the latent factor κm . This quantization ensures that the learned representation aligns with our assumptions about the group characteristics of the latent factors, making it suitable for our objectives. Specifically, to capture the temporal dependency from the sequential observations, sequential neural networks such as Conv1D [52] or Long Short-Term Memory (LSTM) [31] Given that z m produces continuous latent representations, which are incompatible with our requirements, we use a vector quantization layer to discretize the latent space and approximate the latent factor. It maps the continuous representation to the nearest vector in a predefined embedding dictionary E, thereby translating the continuous representations into a discrete latent space. Specifically, the embedding dictionary consists of a set of vectors E = {e 1 , e 2 , . . . , e G }, each representing a distinct group in the discrete embedding space. The assignment of a dictionary vector e i to z m is realized by finding the nearest neighbor in the dictionary as κm = arg min ei ∥z m -e i ∥ 2 , where κm represents the quantized vector that is the closest embedding e i to the continuous representation z m .To effectively estimate latent factors in an unsupervised manner, reconstruction is important because it ensures that the distribution learned by the model closely matches the true observed distribution. According to Definition 4.1, the estimated latent factor thus closely approximates the true latent factor. Given the nature of the transition processes, a conditional decoder is designed with the state-action pairs (s m t-1 , a m t-1 ) as conditions to guide the reconstruction of ŝm t . These conditions, together with the estimated latent factors κm , serve as inputs to the decoder. The accuracy of the reconstruction is quantitatively evaluated by its reconstruction likelihood p Recon (ŝ m t |s m t-1 , a m t-1 , κm ), where p Recon denotes the reconstruction distribution. It provides a probabilistic measure of how accurately ŝm t reconstructs s m t and a quantitative evaluation of the model's reconstruction accuracy.Figure 2: (a) Latent estimation framework takes each trajectory s 0:T as input through a quantized encoder to estimate the latent factor κ. A conditional decoder uses (s t-1 , a t-1 ) as the condition and κ as the input to reconstruct ŝt . (b) After assigning the estimated latent factors to each trajectory, the policy learning framework incorporates the latents as augmented labels to optimize the RL policy. For new individuals, the initial policy is adapted according to their group affiliation, allowing for individualized policy adaptation for newcomers.Training Objectives The parameters are optimized according to the following ELBO objective:EQUATION)where α and β are weights for the corresponding loss components. Specifically, (1) Reconstruction loss L Recon = t ∥s m t -De(En(s m 0:T ), s m t-1 , a m t-1 )∥ 2 , where En and De are the encoder and decoder, measures the discrepancy between the reconstructed state ŝm t and the original state s m t .(2) Quantization loss assesses the discrepancy between the encoder output z m and the discretized representation e m , formulated as L Quant = i ∥sg[z m,i ] -e m,i ∥ 2 . Since the quantization step is undifferentiable, we use the stop-gradient operation sg[•] to enable gradient-based optimization, which updates the dictionary without affecting the encoder parameters. 3) Commitment loss is designed to minimize the discrepancy between z m and e m , ensuring that z m aligns more closely with the embedding space and formulated as L Commit = i ∥z m,isg[e m,i ]∥ 2 . By applying the stop gradient to z m,i , it ensures that the gradients from this loss do not change the dictionary vectors but rather optimize the encoder.The estimation network is pre-trained offline. When a new individual arrives, we estimate κ and adapt the policy simultaneously, and the policy is adapted through new interactions. Specifically, Latent-based Policy Individualization The estimated latent individual-specific factors κ, together with the offline trajectories over all individuals, are used to learn the individualized policy π * κ . We view the estimated factors as an augmented component of the policy input and adjust the policy training objective to match the unique characteristics of each individual. Take Q-learning [58] as an example. In the individualized process, the latent factor is augmented as a policy input represented as µ π (s t ; θ µ ) → µ m π (s m t , κm ; θ µ ), where θ µ represents the parameters of the policy network. The policy model, by incorporating latent factors, can more effectively adapt to the unique characteristics of each individual The training objective is updated accordingly asJ (θ µ ) = E ∞ t=0 γ t Q s t , µ m π (s m t , κm ; θ µ ); θ Q, where γ is the discount factor and Q is the Q value. Such individualization improves the adaptability of the policies to various environments, and our framework is general enough to be integrated with many RL algorithms.Policy Adaptation for New Individual The policy adaptation involves two steps: initializing the policy based on the individualized policy π * κ and fine-tuning the policy through new interactions. For a new individual from group g n , we first estimate its group factor κn and then initialize the policy network π new by directly transferring the parameters from π * κ=κn . Specifically, the new agent updates π new based on the collected trajectory D n by maximizing the expected reward as π new = arg max π E (st,at)∈Dn (R at (s t )). The dataset D n is further augmented with new observations (s t , a t , R t , s t+1 ) from the new individual under the policy π new , which is critical for accurately estimating the latent factor of the new individual. During this process, the policy is iteratively improved to better fit the specific characteristics of the new individual. 
experiment: 
 Evaluation Metrics To measure the latent identification, we quantify the correlation between the estimated and true latent factors by: (1) Pearson Correlation Coefficient (PCC) for single latent, which quantifies the linear correlation between individual estimated and true factors, and (2) Kernel Canonical Correlation Analysis (KCCA) for multiple latents, which evaluates the correlation between sets of estimated and true factors. An absolute value close to 1 indicates a strong correlation and better latent recovery. To evaluate the control performance, we measure the adaptation performance using: (3) jumpstart, which records the improvement in initial performance when a learning agent leverages knowledge from source tasks, and (4) accumulative reward, which indicates the learning quality over the learning process. (5) initial and final reward, which measures the initial performance benefited from policy adaptation and the performance after the full training process.Baselines For estimation evaluation, our baselines include: (1) disentangled sequential autoencoder [88] , which disentangles the latent representations into static and dynamic parts instead of considering the global influence of κ. (2) Population-level component, which embeds the latent factors using population data rather than on an individual basis. For policy evaluation, our baselines include: (3) aligned latent models [22] , which jointly optimizes a latent-space model and a policy to achieve high returns. ( 4) Soft Actor-Critic (SAC) [26] , which uses entropy as part of the objective function to encourage exploration and improve robustness. ( 5) Deep Deterministic Policy Gradient (DDPG) [58] , which combines deterministic policy gradients with deep neural networks to effectively handle continuous action spaces. ( 6) Dueling Double Deep Q-Network (D3QN) [81] , which introduces a dueling architecture for value function estimation and improves value estimation. (7) Rainbow DQN [30] , which combines prioritized experience replay and dueling network architectures to improve performance and learning stability.Synthetic Experiments We first run experiments on the synthetic dataset to demonstrate the effectiveness of the estimation framework, which is manually generated following the post-nonlinear model. We design three types of latent factors κ, each either satisfying or violating the required assumptions. The result shows that satisfying the sample sufficiency assumption is necessary to recover the latent factor. In addition, we evaluate the effect of the trajectory length as described in the Theorem 4.2. The findings, shown as solid lines, indicate that increasing the sample size would apparently improve the identifiability performance, which is consistent with the proposed theorem. Ablation Study The contributions of the different components in the latent estimation framework are reported in Table 1 . We build on the autoencoder framework with a quantization layer and add each component sequentially to the previous module. Incorporating a sequential encoder significantly improves the identifiability, which is important for the accurate recovery of latent factors. In the implementation, we use a noise estimator during optimization to minimize bias and improve identifiability.The results suggest that the noise estimator contributes to fine-tuning the overall performance of the model, allowing for more accurate and reliable recovery of latent individual-specific factors.We further evaluate our framework on the real-world dataset, Per-suasionForGood corpus [80] , which is widely used for analyzing persuasion strategies [66, 8, 90] . It consists of 1017 person-to-person dialogues and 32 personality traits of each participant. In each dialogue, the persuader tries to convince the persuadee to donate to a charity. In the iMDP context, the state refers to the persuadee's response, the action refers to the persuader's utterance, and the reward is the final donation. Since this offline dataset lacks real-time interactions necessary for assessing control performance, we instead use it to identify the latent personality of each individual. We use BERT [11] as the backbone to embed each utterance into a 768-dimensional feature representation and then use an LSTM encoder followed by the quantization layer to recover the latent personalities.The CCA results under different latent dimensions are shown in Figure 4 (a), demonstrating that our method can achieve remarkable performance on the real dataset by appropriately fine-tuning the latent dimensions. Pendulum Pendulum [5] is a continuous control task for RL study with the goal of swinging up and stabilizing in an upright position. The states are the x-y coordinates and angular velocity, and the action is the torque applied to the pendulum. For simplicity, we choose DDPG as the foundational optimization algorithm and manually create 20 individualized environments. In these environments, the gravity g is randomly drawn from a categorical distribution over the set {3, . . . , 12}.The performance of the policy adaptation is evaluated on a new individual with g = 10.We compare our method against several baselines: (1) SAC; (2) DDPG without prior knowledge; (3) aligned latent models; (4) pre-trained DDPG incorporating knowledge from given individuals, termed cross-individual transfer; (5) individualized policy incorporating randomly defined group embedding, termed random group embedding. The training curves over accumulative reward are reported in Figure 4 (b), showing that the proposed method outperforms other baselines in both jumpstart and accumulative reward. Specifically, methods that benefit from population knowledge (our method and cross-individual transfer) outperform non-transfer methods, indicating that the pre-trained policy would accelerate the learning process. However, since cross-individual transfer ignores the individualspecific information, such mixed policy knowledge yields worse initial performance compared to the individualized policies derived from our method.HeartPole HeartPole [59] is a discrete healthcare environment that explores the long-term health outcomes of short-term decisions. Half Cheetah Half Cheetah [74] is a Mujoco-based task aiming to control a 2D bipedal robot. The agent consists of 9 links and 8 joints, and the goal is to apply torque to the joints to make the cheetah run forward as fast as possible. We introduced 50 individualized settings with the gravity g following a categorical distribution with probabilities p = 0.2 and corresponding g values from {8, 8.5, . . . , 10}. The adaptation performance is evaluated on a new individual with g = 9.8. We compare our method with (1) DDPG without prior knowledge, (2) aligned latent models, (3) crossindividual transfer with DDPG, and (4) random group embedding. The comparative results are shown in Figure 4 (d). We found that inappropriate source domain data can degrade the control performance (see cross-individual transfer), but the integration of group embedding facilitates generalization, allowing our method to outperform baselines in terms of convergence speed and efficiency. 
conclusion and limitations: 
 Our work focuses on learning latent state transitions from observed state-action trajectories, ensuring identifiability even in the presence of latent individual-specific factors. To the best of our knowledge, this study provides novel identifiability guarantees in several settings that have not been addressed by others. Despite these contributions, our approach has three major limitations. (1) It currently does not account for instantaneous causal influences within s t . However, this problem could be mitigated by adjusting the temporal resolution of the data and explicitly modeling causal dependencies within states. Integrating causal graphical models or advanced inference techniques for handling instantaneous causal relationships would enhance our framework. (2) It lacks a nonparametric proof for scenarios where latent factors are continuous, while our empirical results suggest that the approach may be adaptable to these more general conditions. (3) The proposed model does not account for time-varying latent factors. Establishing theoretical identifiability is highly non-trivial and further constraints would be needed.These limitations present key areas for future research. In addition, practical concerns such as privacy, robustness, and reliability are essential for real-world applications. For privacy, de-identification techniques such as removing direct identifiers (e.g., names, postal codes), applying data perturbation, and pseudonymization could help mitigate risks. The exploration of differential privacy techniques is also a promising direction for ensuring privacy and security in practical applications. s = [s 1 , s 2 , . . . , s ds ] ⊤ d s -dimensional observed states a = [a 1 , a 2 , . . . , a da ] ⊤ d a -dimensional observed actions κ = [κ 1 , κ 2 , . . . , κ dκ ] ⊤ d κ -dimensional latent individual-specific factors 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)? Answer: [NA] Justification: Not related to our work. Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: All are referred. Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset's creators. 
new assets: 
 Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? 
 
 Appendix: Given the identifiability theorems, we first provide intuitive explanations for each assumption and discuss their relevance to real-world applications. Then, we provide the proof. Finally, we introduce some preliminaries related to our theorems, which are essential for the proof.B.1 Preliminaries for Theorem 4.1 and 4.2The first-order Markov property implies that the transition probability to the next state depends only on the current state, uninfluenced by the sequence of previous states. Specifically, Definition B.1 (First-order Markov Property [72] ). A stochastic process {X t : t ∈ N } has the firstorder Markov property if, for each set of times t, t-1, . . . , 0 and corresponding state x t , x t-1 , . . . , x 0 in the state space, the following conditional independence property holds:EQUATIONThe first-order Markov property implies that the transition probability to the next state depends only on the current state, uninfluenced by the previous states. In the context of the state transition process, it possesses the first-order Markov property. Mathematically, it can be represented as:EQUATIONwhere P(s t |s t-1 , a t-1 ) is the transition probability from (s t-1 , a t-1 ) to the state s t .A finite mixture model is used for modeling a total population that comprises unobserved or hidden groups. Each of these groups is assumed to follow its own distinct probability distribution. In this context, the overall population model is expressed as a weighted sum of these individual distributions [63] . Specifically, Definition B.2 (Finite Mixture Models [77] ). A finite mixture model is a probability law based on a finite number of probability measures, µ 1 , . . . , µ m , and a discrete distribution ω 1 , . . . , ω m . A realization of a mixture model is generated by generating a component at random k, 1 ≤ k ≤ m, and then drawing from µ k ∼ P. Then, the mixture measure P is defined as a weighted sum of probability measures µ i with weights w i . Specifically,EQUATIONGroup Determinacy The distinct values of the latent factor κ categorize the population into separate groups, with each group characterized by its unique probability distribution and denoted as G g=1 π g δ κg (κ). The mixture formulation implies that the latent factor κ serves as a categorical variable, with each unique value explicitly specifying a distinct group within the population. Such a formulation facilitates the identification and analysis of heterogeneous subpopulations within the finite mixture model.The idea of group determinacy is important in real-world applications. Take personalized education as an example. For each student m, s m represents their current knowledge state, a m denotes their personalized learning action, and the function f determines the unique educational trajectory for each student. The latent factor κ m influences how a student's learning progresses over time. It can be based on factors such as learning style preferences that help to logically group students. Specifically, one group might consist of visual learners who excel in interactive, graphically-oriented subjects, while another group might include students who prefer textual information and excel in reading and writing-intensive subjects. Each group exhibits its own set of learning outcomes and patterns, allowing educators to personalize teaching methods and materials to effectively meet the different needs of each group.In a finite mixture model with G groups, each group requires sufficient observations to identify the latent group factor κ. This assumption provides a minimum number of observation samples, which is 2G -1 observations in each group. Such a threshold ensures that we have enough information and variability in the observed data to distinguish the characteristics of each group. This assumption helps to identify the unique characteristics of each individual, which is critical for identifiability.Sample sufficiency indicates that sufficient data are needed to achieve identifiability, which is a fundamental assumption in many analytical models. For example, in the context of nonlinear ICA using auxiliary variables [39] , it is necessary to have at least 2n + 1 values for the auxiliary variables to ensure sufficient variability and guarantee identifiability. Similarly, for successful disentanglement with minimal change [47] , at least 2n + 1 domain embeddings are required to ensure identifiability. Intuitively, without sufficient data to provide us with relevant information about the parameters, it is impossible to determine the values of these parameters.Infinite Samples and Overlapping Conditions Asymptotic identifiability refers to the property that a model becomes identifiable as the sample size goes to infinity. In practical terms, this means that given an infinite amount of data, one would be able to consistently estimate the parameters of the model. In the context of finite mixture models, if there are enough samples for each individual, then the corresponding components can be identified directly from each individual [37] .The overlapping condition requires the existence of at least two different individuals within the same group of the population, who have identical conditional probabilities. This assumption is crucial as it ensures that the model accounts for overlapping behavioral responses between different individuals, which is a common phenomenon in heterogeneous populations. Consider personalized education as an example, where students come from different academic backgrounds and have different levels of prior knowledge. Despite this initial heterogeneity, it is possible for two students in the same learning group to have the same probability of successfully completing a task.We first show that the individualized transition processes can be viewed as a finite mixture model with grouped samples and then derive the identifiability under two scenarios.Lemma B.1 addresses the identifiability of mixture models from grouped samples. It implies that with sufficient data per group, each component of the mixture model can be determined without ambiguity from the observed data. Specifically, suppose we have a mixture model consisting of G different probability distributions, that is, G components c 1 , c 2 , . . . , c G . Each component c i corresponds to a unique probability density function P i (•). These components are mixed together, with each component having a mixing weight π i , satisfying π i ≥ 0 and G i=1 π i = 1. Now suppose we have G observation groups g 1 , g 2 , . . . , g G , with each group g i containing observations that are independent and identically distributed drawn from the same component c i . Lemma B.1 (Identifiability of Mixture Models from Grouped Samples [77] ). Suppose we have observations from a mixture model and that they are grouped such that observations in the same group are known to be drawn from the same component. Denote by G the number of groups. If there are at least 2G -1 observations per group, any mixture of G probability measures can be uniquely identified.Proof. The observation model for each group g i with a unique latent factor κ i can be defined as P(s, a, s ′ |κ i ) = P(s, a, s ′ , u|κ i )du, where u denotes individual-specific factors within the group g i . Since u comprises both κ and the initial state z, we have P(u) = P(u)P(κ|u) = P(u, κ) = P(u|κ)P(κ). In this work, we consider two scenarios to provide identifiability from the observed temporal data collected from the population. Theorem 4.1 Assumptions in theorem 4.1 and Eq. 2 ensure that individuals within the same group share identical joint distributions. Suppose the observations can be grouped into G finite components, then the joint distribution can be factorized as: P(s, a, s ′ ) = P(s, a, s ′ |u)P(u)du= P(u)P(s ′ |s, a, κ)P(s, a|u)du= P(u|κ)P(κ)P(s ′ |s, a, κ)P(s, a|u)dudκ= P(κ)P(s ′ |s, a, κ)P(s, a|u)P(u|κ)dudκ= P(κ)P(s ′ |s, a, κ)P(s, a, u|κ)dudκ= P(κ)P(s ′ |s, a, κ)P(s, a|κ)dκ= P(κ)P(s ′ , s, a|κ)dκ (14) This formulation asserts that the joint distribution of P(s, a, s ′ ) for the entire population can be modeled as a mixture model governed by the respective κ i values. Since sample sufficiency ensures that the sample size of observations within each group is greater than 2G -1, then the identifiability of κ is guaranteed by Lemma B.1.Theorem 4.2 According to assumptions in theorem 4.2, if each individual's trajectory is sufficiently long, the conditional model P(s ′ |s, a), as well as P(s, a) for the individual, would become asymptotically identifiable. Consider an extreme scenario where each individual is treated as a distinct group. Asymptotically, it is possible to identify the mixture distribution of the population. However, some individuals may share the same κ and can be grouped together. Then, we need to find the similarity between different individuals and merge them into the same group.An intuitive merging criterion is as follows. Asymptotically, it can be inferred that for some particular value of (s, a) = (s * , a * ), the probability of s ′ given (s * , a * ) will be the same across individuals in the same group. Define t j as the time of the j-th occurrence of (s * , a * ). We can then define the collection of variables X j as X j = {s t j +1 , t j = 1, . . .}, representing the state at time t + 1 given the fixed state and action (s * , a * ) at time t. In this way, for any j, X j sampled from a particular group j.Then, the identifiability is ensured by Lemma B.1.Remark B.1. Prior work [37] used a Gaussian mixture model as a prior on the coefficients, while the latent confounder variable, denoted as Z, was constrained to a binary state, thereby indicating group membership for a given individual. We extend this foundation by generalizing the latent confounder to a set of discrete values and considering a nonparametric model for broader applications.Recent work [12, 36] provides the necessary and sufficient conditions for the identifiability of certain latent structural patterns, but it rules out the case of triangle structure involving latent variables. In this paper, we extend their work to the temporal case and provide the identifiability of the latent group factor κ, where κ can be either continuous or discrete. Furthermore, we allow multiple instances of κ to influence the state transition dynamics.In Theorem 4.3, we aim to identify the latent group factors based on a post-nonlinear temporal causal model, as shown in Figure 5 . We assume the existence of a learnable and invertible embedding mapping f , which is able to preserve the causal structure intrinsic to the state s. Specifically, Definition C.1 (Post-nonlinear Temporal Causal Models). Consider a scenario where d s observed states from the k-th individual are denoted as s k t = (s k 1,t , . . . , s k ds,t ) T , which is a direct observation of an embedded representation f (s), alongside m unobserved group factors κ = (κ 1 , . . . , κ m ) T . The state transition dynamics satisfyEQUATION)for i = 1, . . . , n. Here, α ij and β ij represent causal coefficients that quantify the influence of the state f (s j,t ) and the action a j,t on f (s i,t+1 ), respectively. P i and L i denote the sets of direct state and action that influences s k i,t+1 (or f (s k i,t+1 )). Actions are considered to be stochastic. The coefficients λ j are individual-specific and show variation across individuals. The random noise term ϵ k i,t+1 is independent of s j,t and a j,t for all j ∈ N + to account for unmeasured influences. Let S t = {s 1,t , s 2,t , . . . , s ds,t } represent the set of all state variables in the system at any time t = 1, . . . , T , and let L = {κ 1 , κ 2 , . . . , κ m } represent the set of latent variables, where m and d s are the numbers of latent and state dimensions, respectively. A subset R t,t -⊆ S t ∪ S <t (or R t,t + ⊆ S t ∪ S >t ) with cardinality r, is called a minimal rank set if it satisfies the following conditions:(i) The bottleneck set, defined as B = L ∪ S t for any given time t, can t-separate (see Definition C.5) any pair of minimal rank sets (R t,t -, R t,t + ), where t -< t < t + .(ii) There does not exist a subset R ′ t,t ± ⊂ R t,t ± with |R ′ t,t ± | < |R t,t ± | that can satisfy condition (i).The set R t,t ± is considered minimal in the sense that it is the smallest cardinality subset of observed state variables that includes a bottleneck set and disjoint state variables, capable of representing the essential separation status within the system. An illustrative example of a minimal rank set (see Figure 6 ) is shown in the yellow area, and a bottleneck set is depicted in the green area.C. EQUATIONand equality holds for generic covariance matrices consistent with G. Lemma C.2 (Identifiability of Linear Regression Models). Consider a linear regression model with a response variable Y and p predictors X 1 , X 2 , . . . , X p . The linear relationship is defined as:EQUATION)where β 0 , β 1 , . . . , β p are the regression coefficients and ε is the error term. The matrix representation can be expressed as:EQUATION)where Y is the response vector, X is the design matrix including predictors, β is the vector of regression coefficients, and ε is the vector of error terms. For the regression coefficients β to be identifiable, the design matrix X must have full column rank, meaning no predictor is a perfect linear combination of the others. This ensures that the matrix X T X is invertible, allowing for the unique estimation of β through:β = (X T X) -1 X T Y (19) Lemma C.Proof. Suppose latent factors exist and influence the embedding of the observed states f (s), which preserve the causal structure intrinsic to the states s. According to Lemma C.1, the rank of Σ A i ,B i should be less than or equal to min{|CA i | + |C B i |}.In the absence of latent factors, according to the CI skeleton, the minimal configuration to t-separate A i from B i is byEQUATIONConsequently, the rank of the covariance matrix is rank(Σ A i ,B i ) = |{f (s 1,i ), . . . , f (s ds,i )}| + |∅| = d s .If the calculated rank is greater than d s , it implies the presence of latent variables accounting for the unexplained variance since the observed variables alone would not result in such rank deficiency.In scenarios with latent factors, the maximum rank deficiency observed across covariance submatrices, representing the discrepancy between the expected and actual ranks, establishes a lower bound for the number of latent variables. Considering the minimal t-separation of A i and B i occurs via ({f (s 1,i ), . . . , f (s ds,i ), κ 1 , . . . , κ m }, ∅) or (∅, {f (s 1,i ), . . . , f (s ds,i ), κ 1 , . . . , κ m }).Then the rank of the covariance matrix is rank(Σ A i ,B i ) = |{f (s 1,i ), . . . , f (s ds,i ), κ 1 , . . . , κ m }| + |∅| = m + d s . By iteratively computing the rank of rank(Σ A i ,B i ), a consistent value corroborates the existence of latent factors influencing all observed states. Furthermore, the count of latent factors can be deduced by m = rank(Σ A i ,B i ) -d s .In conclusion, under the conditions of the theorem, if the observed rank deficiency in the covariance matrix of observed variables cannot be explained by the observed variables alone, it implies the existence of latent variables. Furthermore, the number of such latent variables can be inferred from the extent of the rank deficiency.Proof. For each individual k, consider the proposed model at any time t and t + 1:f (s k i,t ) = j∈Pi α ij f (s k j,t-1 ) + j∈Li β ij a k j,t-1 + m j=1 λ j κ j + ϵ k i,t , f (s k i,t+1 ) = j∈Pi α ij f (s k j,t ) + j∈Li β ij a k j,t + m j=1 λ j κ j + ϵ k i,t+1 .Subtracting these two equations, we obtain:f (s k i,t+1 ) -f (s k i,t ) = j∈Pi α ij (f (s k j,t ) -f (s k j,t-1 )) + j∈Li β ij (a k j,t -a k j,t-1 ) + (ϵ k i,t+1 -ϵ k i,t ). Define x k i,t+1 = f (s k i,t+1 ) -f (s k i,t ), y k i,t+1 = a k i,t+1 -a k i,t, and η k i,t+1 = ϵ k i,t+1 -ϵ k i,t . Substituting these, the model transforms to:x k i,t+1 = j∈Pi α ij x k j,t + j∈Li β ij y k j,t + η k i,t+1 .In that case, the identifiability of α and β can be directly derived by Lemma C.2. We further assume that the m latent factors follow the Normal distribution. Drawing on methodologies used in factor analysis C.3, then λ is orthogonal-wise identifiable.In this section we present four illustrative examples to describe cases where identifiability is achieved.For the sake of simplicity, we define X as X = f (s) and omit the terms a and ϵ from our illustration for simplicity, under the assumption that they are random and independent variables.In Example 1, as shown in Figure 6 , there is only one latent factor and onedimensional states. The bottleneck set is (X 3 , κ), and the pairs of minimal rank sets are (A, B) = ((X 1 , X 2 , X 3 ), (X 3 , X 4 , X 5 )). According to the causal graph, the minimal way to t-separate A from B is either ({κ, X 3 }, ∅) or (∅, {κ, X 3 }). Consequently, the rank of the covariance matrix is rank(Σ A,B ) = |{κ, X 3 }| + |∅| = 2. According to Theorem 2, the fact that rank(Σ A,B ) = 2 > 1 indicates the presence of the latent factor. Consequently, the number of latent variables can be deduced as m = rank(Σ A,B ) -1 = 2 -1 = 1. Example 3 As shown in Figure 8 , there is one latent factor and two-dimensional states. The bottleneck sets are ((X 13 , X 23 , κ), (X 14 , X 24 , κ)), and one possible pairs of minimal rank sets are (A 1 , B 1 ) = ((X 11 , X 12 , X 13 , X 23 ), (X 13 , X 23 , X 24 , X 25 )) and (A 2 , B 2 ) = ((X 12 , X 13 , X 14 , X 24 ), (X 14 , X 24 , X 25 , X 26 )). According to the causal graph, the minimal configuration to t-separate A 1 from B 1 is either ({κ, X 13 , X 23 }, ∅) or (∅, {κ, X 13 , X 23 }). Consequently, the rank of the covariance matrix is rank(Σ A1,B1 ) = |{κ, X 13 , X 23 }| + |∅| = 3. Similarly, the minimal configuration to t-separate A 2 from B 2 is either ({κ, X 14 , X 24 }, ∅) or (∅, {κ, X 14 , X 24 }), resulting in rank(Σ A2,B2 ) = |{κ, X 14 , X 24 }| + |∅| = 3. According to Theorem 2, the fact that rank(Σ A1,B1 ) = rank(Σ A2,B2 ) = 3 > 2 indicates the presence of the latent factor. Consequently, the number of latent variables can be deduced asm = rank(Σ A i ,B i ) -2 = 3 -2 = 1.Example 4 As shown in Figure 9 , there are two latent factors and two-dimensional states. The bottleneck sets are ((X 14 , X 24 , κ 1 , κ 2 ), (X 15 , X 25 , κ 1 , κ 2 ), and one possible pairs of minimal rank sets are (A 1 , B 1 ) = ((X 11 , X 12 , X 13 , X 14 , X 24 ), (X 14 , X 24 , X 25 , X 26 , X 27 )) and (A 2 , B 2 ) = ((X 12 , X 13 , X 14 , X 15 , X 25 ), (X 15 , X 25 , X 26 , X 27 , X 28 )). According to the graph, the minimal configuration to t-separate A 1 from B 1 is either ({κ 1 , κ 2 , X 14 , X 24 }, ∅) or (∅, {κ 1 , κ 2 , X 14 , X 24 }).Consequently, the rank of the covariance matrix is rank(Σ A1,B1 ) = |{κ 1 , κ 2 , X 14 , X 24 }| + |∅| = 4.Similarly, the minimal configuration to t-separate In this discussion, we introduce the concept of the covariance matrix within the framework of latent variable models. By examining the properties of the covariance matrix, such as its rank, we are able to identify signs of latent variables-rank deficiencies, which serve as a measure of the cardinality of the minimal set of latent variables required to explain the observed dependencies. Such rank deficiencies indicate the presence of latent variables that extend beyond the observable scope.A 2 from B 2 is either ({κ 1 , κ 2 , X 15 , X 25 }, ∅) or (∅, {κ 1 , κ 2 , X 15 , X 25 }) with rank(Σ A2,B2 ) = |{κ 1 , κ 2 ,Consider a directed acyclic graph (DAG), denoted as G, whose vertices V (G) form the set [m] := {1, 2, . . . , m}. Each node i in G is associated with a random variable X i and an independent error term ϵ i ∼ N (0, ϕ i ) with ϕ i > 0. The DAG structure imposes a recursive relationship among the variables, where the value of X j can be expressed as a linear combination of the variables X i of its parent vertices pa(j), alongside the error term ϵ j and regression coefficients λ ij that correspond to the edges i → j in G:EQUATION)where pa(j) denotes the set of parent nodes of vertex j, where a parent node i is one that has an edge leading to j in G. From this recursive sequence of regressions, one can solve for the covariance matrix Σ of the jointly normal random vector X, which is defined as follows.Definition C.3 (Covariance Matrix of Random Vector [70] ). The covariance matrix of the random vector is given by the matrix factorizationEQUATION)where matrix Φ is defined as a diagonal matrix with the variances of the error terms as its diagonal elements: Φ = diag(ϕ 1 , . . . , ϕ m ). The matrix M is an m × m upper triangular matrix where M ij = λ ij if i → j isThe concepts of Trek and Trek Separation precede a crucial need to address the presence of latent variables and intricate dependency structures that are not directly observable. The Trek represents a particular path that interconnects variables within a graph, even if they are not directly linked, while Trek Separation delineates a criterion to ascertain whether two sets of variables are independent, conditional on a set of other variables. Below, we give the formation definitions of these two concepts. Definition C.4 (Trek [70] ). A trek in G from i to j is an ordered pair of directed paths (P 1 , P 2 ) where P 1 has sink i, P 2 has sink j, and both P 1 and P 2 have the same source k. The common source k is called the top of the trek, denoted top(P 1 , P 2 ). Note that one or both of P 1 and P 2 may consist of a single vertex, that is, a path with no edges. A trek (P 1 , P 2 ) is simple if the only common vertex among P 1 and P 2 is the common source top(P 1 , P 2 ). We let T (i, j) and S(i, j) denote the sets of all treks and all simple treks from i to j, respectively. Definition C.5 (Trek Separation [70] ). Let A, B, C A and C B be four subsets of V (G) which need not be disjoint. We say that the pair (C A , C B ) trek separates (or t-separates) A from B if for every trek (P 1 , P 2 ) from a vertex in A to a vertex in B, either P 1 contains a vertex in C A or P 2 contains a vertex in C B .The Conditional Independence (CI) skeleton in graphical models refers to a structure that represents the conditional independence among observed variables. The CI skeleton can be used to infer the existence of latent variables. If the observed data suggests dependencies not represented in the CI skeleton, it may indicate hidden factors at play. The formal definition is given as follows. Definition C.6 (Conditional Independence Skeleton [12] ). A CI skeleton of X is an undirected graph where the edge between X 1 and X 2 exists if and only if there does not exist a set of observed variablesC such that X 1 , X 2 / ∈ C and X 1 ⊥ ⊥ X 2 |C.Reinforcement Learning In RL, an agent learns to make decisions by interacting with the environment. The agent receives rewards for taking actions in the environment and uses this feedback to learn optimal behavior. It is often modeled as a Markov Decision Process (MDP) represented by a tuple S, A, P, R, γ , where S denotes a finite set of states representing different situations an agent might encounter, A a finite set of actions representing different decisions an agent can make, P a state transition function defining the probability of transitioning to a new state s ′ given a current state s and action a, denoted as P(s ′ |s, a), R a reward function assigning a scalar value to each state-action pair (s, a), representing the immediate reward received after performing action a in state s. γ ∈ [0, 1] is the discount factor, representing the agent's consideration for future rewards. The agent's goal is to learn an optimal policy π * , which defines the optimal set of actions in different states to maximize the expected cumulative discounted reward over the long run. Developing this optimal policy involves estimating value functions such as the action-value function, defined asQ π (s, a) = E π ∞ t=0 γ t R t |S 0 = s, A 0 = a ,which represents the expected reward of taking action a in state s following policy π. The pursuit of optimal policy π * involves maximizing the value functions over all possible state-action pairs:π * = arg max π Q π (s, a).Variational Autoencoder Variational Autoencoders (VAEs) [44] 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 ϕ (z|x) and a decoder p θ (x|z). The encoder maps input data x to a latent space, represented by a probability distribution, typically Gaussian, with parameters µ and σ signifying the mean and standard deviation, respectively. The decoder reconstructs the input data from a sampled latent representation z.The distinct feature of VAEs lies in their probabilistic approach. The encoder outputs parameters of a latent distribution, from which a sample z is drawn:EQUATIONThe decoder then attempts to reconstruct the input from this latent sample. VAEs optimize the Evidence Lower Bound (ELBO) objective, which balances two aspects: the reconstruction quality and the regularization of the latent space. The traditional ELBO is given by:EQUATIONHere, the first term measures the reconstruction quality, while the second term, the Kullback-Leibler (KL) divergence, imposes a regularization by encouraging the latent distribution q ϕ (z|x) to be close to a prior p(z), typically assumed to be a standard normal distribution N (0, I). VAEs, through this optimization, are capable of generating new data points that are similar to the input data, making them highly valuable in applications like image generation, denoising, and anomaly detection within the domain of unsupervised learning.Individualized Machine-Learning Applications In the modern era, the power of machine learning has been harnessed to create highly individualized solutions across a myriad of domains. In the realm of health and wellness, machine learning aids in tailoring interventions for increasing physical activity [89, 57] , promoting weight loss [18, 17] , improving adherence for diabetes [89] . For the elderly, personalized algorithms assist in both technology adaptation and specialized care for conditions [32] .The financial sector benefits from machine learning's prowess in optimizing technical indicators, making stock market predictions more precise and individualized [56] . In the educational landscape, Information and Communication Technology (ICT) leverages machine learning to offer personalized education systems such as adaptive e-learning interfaces [16] and individualized tutorial planning [40] . Furthermore, the transportation sector sees advancements with car-following control strategies tailored for individual drivers [69] . Multimedia platforms, such as YouTube and TikTok, are enhancing user experiences by offering video content recommendations fine-tuned to individual preferences using reinforcement learning [6, 33] . These examples merely scratch the surface, emphasizing the vast and diverse applications of individualized machine learning in today's world.Reinforcement Learning for Latent State-Transition Processes RL has witnessed significant advancements in recent years, particularly with the integration of latent variable models to capture the underlying dynamics of environments. A primary focus in this domain is learning low-dimensional, latent Markovian representations from observed data [54, 48, 42, 25, 82, 95, 49, 62, 19, 23, 20, 91] .Common strategies for state representation learning include reconstructing the observation, learning a forward model, or learning an inverse model. Additionally, prior knowledge, such as temporal continuity [83] , can be leveraged to constrain the state space. Numerous studies have proposed methods to estimate the underlying state-transition process from high-dimensional input sequences [82, 13, 25, 27, 94, 19, 41, 28] . Using the learned world model, agents can engage in model-based RL or planning. Furthermore, these methods encode structural constraints, ensuring the sufficiency and minimality of the estimated state representations from both generative and selection processes.Recently, several studies [60, 55, 78, 79, 4, 65] have aimed to estimate the state-transition process in the presence of latent confounders. A handful of work [60, 65] can be viewed as addressing similar settings involving individual-specific factors. However, to the best of our knowledge, we have yet to identify a systemic approach that offers a clear identifiability result for the state-transition process when individual-specific factors are present.Comparisons with Related Works Contextual MDPs [29] consider the general contextual influence on transition probabilities and rewards. However, the context variables are assumed to be partially observable and do not guarantee the identifiability of the context variables. In our case, when the latent factors are finite, our method guarantees group-wise identifiability even when the transition processes are nonparametric. In the cases of infinite latent factors, identification could be achieved under proper assumptions.Multi-task RL [75] involves learning policies for a variety of tasks simultaneously. The goal of the agent is to perform well on all these tasks, which may have similar or different objectives. often involves sharing information between tasks to improve learning efficiency and policy performance. Instead of focusing on policy optimization for all tasks, our work identifies latent individual-specific factors that implicitly influence the decision-making process. These factors indicate the unique properties of each individual, providing explanatory clues for policy adaptation.Meta-RL [3] trains a learning model on a variety of tasks so that it can efficiently apply what it has learned to new tasks. Unlike our method, it does not assume a time-invariant latent factor, has no guarantee of identifiability, and does not provide a clear clue of adaptation. While iMDP captures how an individual's belonging to a certain group affects their interactions within an environment, allowing for individualized policy adaptation. Moreover, iMDP provides a guarantee of identifiability and develops a corresponding estimation framework that potentially offers better interpretability.Factored MDP [92] and Factored Non-stationary MDP [24] assume there are no unobserved confounders in the state transitions. Block MDP [10] , POMDP [38] , and Latent MDP [86] consider latent states/spaces, but such latent factors do not influence each state in the transition process. Specifically, existing works [92, 10, 24] 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 [38, 86] . However, the identifiability of time-invariant latent confounders, though not well-studied, has numerous applications. The aforementioned settings differ significantly from our work since we consider a latent group 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 the case of nonparametric transitions. This establishes novel theoretical insights for learning state-transition processes with latent factors.Pearson Correlation Coefficient Pearson Correlation Coefficient (PCC) [9] is a statistical measure that quantifies the degree of linear relationship between two variables. It provides a value between -1 and 1, where 1 implies a perfect positive linear relationship, -1 implies a perfect negative linear relationship, and 0 implies no linear relationship between the variables. The equation for calculating the Pearson Correlation Coefficient r between two variables X and Y is as follows:EQUATION)where n is the number of paired samples, xy is the sum of the product of paired scores, x and y are the sums of the x scores and y scores respectively, x 2 and y 2 are the sums of the squared x scores and y scores respectively.Canonical Correlation Analysis Canonical Correlation Analysis (CCA) [34] is designed to identify bases for two sets of variables in order to maximize the mutual correlations between the projections onto these bases. In our work, CCA is used as an evaluation metric to validate that the recovered latent variable is meaningfully related to the ground truth latent variable, thus proving the relevance of the estimated representations. Let X and Y be the two sets of observed variables. This algorithm starts by centering the columns of X and Y so that they have zero mean. Then the covariance matrices C XX = X ⊤ X, C Y Y = Y ⊤ Y , and C XY = X ⊤ Y are calculated. After that, the canonical correlations are obtained by solving the following generalized eigenvalue problem:C -1 XX C XY C -1 Y Y C Y X ν = λν.The square roots of the eigenvalues λ indicate the canonical correlations between the linear combinations of X and Y . The corresponding eigenvectors ν and u = C XY ν are the canonical weights used to construct the canonical variables. Finally, the canonical variables of X and Y are U = Xν and V = Y u, respectively, representing the linear combinations of the original variables that are maximally correlated. The correlation of the primary pair of canonical variables is the highest, followed by the secondary pair, and so on. When employing CCA as an evaluation metric, a higher canonical correlation indicates a stronger and more relevant relationship between the recovered latent variable and the ground truth latent variable.To extend the capability of CCA for analyzing nonlinear relationships, Kernel Canonical Correlation Analysis (KCCA) [84] is employed in the experiment which uses kernel functions to map the original variables into a higher-dimensional feature space. This allows for the capture of more complex, nonlinear correlations between the variables, thus potentially increasing the robustness and relevance of the relationships discovered in scenarios where linear methods fall short.Synthetic Data Generation Processes In this paper, we created three synthetic datasets: Case 1 corresponds to a finite latent factor that satisfies our assumptions, and Case 2 and Case 3 allow for multiple finite and infinite latent variables. The dimensions of states and actions are set to 3 and 2, respectively. The actions taken are generated randomly, following a uniform distribution Uniform(0, 1). The noise term follows a mean-zero Gaussian distribution. The mixing function f corresponds to the post-nonlinear model [93] , where f 1 represents the nonlinear effect, and f 2 denotes the invertible post-nonlinear distortion on s t , embodied by a randomly initialized three-layer MLP with Tanh activation function. The data generation process follows:EQUATION)PersuasionForGood The PersuasionForGood dataset reveals the mechanics of persuasion in the context of charitable giving. It contains 1017 dialogues from 1285 participants in which one participant, called the persuader (ER), tries to convince the other participant, called the persuadee (EE), to donate to a charity. An example dialog is shown in Figure 11 . All participants underwent personality assessments, which included detailed participant-level information such as demographics, Big Five personality traits, moral foundations, and so on, allowing for a multifaceted analysis of persuasion strategies and allowing us to use the labeled 32-dimensional personalities of each persuader as the ground-truth latent factor in our experiments.We use this dataset to evaluate the performance of our estimation framework. We compare the estimated factors to the documented personality traits of persuaders. By examining the interactions between participants with different backgrounds and personalities, we aim to identify underlying patterns that could create effective persuasive agents. Specifically, we use BERT embeddings to generate a 768-dimensional feature vector for each dialog utterance. This process starts with tokenization, segmenting words into smaller units. BERT then processes these tokens to produce contextual embeddings. Pendulum The pendulum environment, provided by OpenAI Gym, is a classic control task used for the evaluation RL models. This environment presents a continuous control task where the agent must learn to control a frictionless pendulum with the goal of swinging it to the highest point and keeping it in the inverted position. The pendulum starts at a random position, and the goal is to bring it to a standstill at the inverted position with the least amount of effort. The system is characterized by a continuous action space, representing the torque applied to the pendulum's fulcrum. For a pendulum of length l and mass m, subject to gravity g and a control input u, the equations of motion can be described by the following second-order nonlinear ordinary differential equations θ = ω, ω = -g l sin(θ) + u ml 2 , where θ is the angle of the pendulum from the vertical upright position, and ω is the angular velocity of the pendulum. The state of the pendulum at any time t can be represented as ⃗ s t = [cos θ t , sin θ t , ω t ], action represents the torque applied to the free end of the pendulum in the range a t ∈ [-2, 2], and the reward function is defined as: r t = -(θ 2 t + 0.1 * ω 2 t + 0.001 * a 2 t ). The goal of RL algorithms is to determine an optimal control policy π * that minimizes the effort to swing and balance the pendulum upright, typically by minimizing a cost function defined over states and actions. Each episode provides a continuous stream of observations, actions, and rewards, allowing the development and evaluation of algorithms capable of learning effective control policies in continuous action spaces. In academic studies, the Pendulum environment serves as a benchmark to investigate the effectiveness of RL algorithms in handling continuous control tasks.HeartPole HeartPole provides a straightforward scenario for assessing healthcare treatment, highlighting the complex interplay between productivity, health, and decision-making strategies. It simulates a professional's quest for increased productivity and examines the long-term health impacts of short-term choices, such as insufficient sleep, and intake of coffee and alcohol. The states include alertness, hypertension, intoxication, time since last sleep, total elapsed time, and total work done.A productivity function and a heart attack risk function are defined over these variables, rewarding incremental productivity while imposing a significant penalty for heart attacks. Every thirty minutes, the agent evaluates the current state and chooses from a set of actions: work, drink coffee (which increases alertness and hypertension), drink alcohol (which decreases alertness while increasing hypertension and intoxication), or sleep (time-consuming but essential to reduce hypertension and intoxication to maintain alertness).Half Cheetah Half Cheetah is an integral part of the Mujoco physics engine, designed to simulate the agility and mechanics of a cheetah through a 2D robotic model. This model consists of 9 body parts, including a torso, two front and two back thighs, shins and feet, connected by 8 joints to provide fluid motion reminiscent of a cheetah's natural gait. The primary goal for this robot is to achieve maximum forward speed while maintaining stability, mirroring the efficiency and speed of its biological counterpart.The observational data in this environment includes both the position and velocity of each segment of the half cheetah, methodically ordered with all position data provided before velocity information. This systematic ordering allows for a detailed understanding of the dynamics of the robot at any given time. Actions within this simulation are defined by the torque applied to the joints, which directly affects its acceleration and motion patterns. The reward function for Half Cheetah is designed to encourage rapid forward motion and operational efficiency and consists of two main components: a forward motion reward proportional to the increase in the robot's horizontal displacement over time and a control cost penalty for unnecessary control effort and external force application. This reward structure is carefully designed to encourage the optimization of forward motion, with an emphasis on reducing control effort and mitigating forces that may interfere with the robot's streamlined motion.Ablation Study: Variability in number of samples Here, we vary the number of samples to further verify this effectiveness. The data generation process is the same as Case 1, except that we change the number of samples to {100, 150, 200, 300, 500, 800, 1000}. The comparison results shown in Figure 14 indicate that our method can achieve consistently good recovery performance under different numbers of individuals. This further confirms that the identifiability of our framework is guaranteed by the mathematical relationship between the trajectory length and the number of groups, which is constrained by the sample sufficiency assumption under the conditions given in Theorem 4.1. Moreover, Figures 14(b ) and 14(c) show the successful recovery of the latent group factor, validated by high-frequency similarity and a remarkable PCC value, confirming the ability of our method to skillfully recover latent variables in practical pendulum tasks.Ablation Study: Variability in initial states We conduct additional experiments to verify how the variability in initial states across individuals affects performance. The initial state distributions are defined with two types: normal and uniform. For the normal distribution, the means are set to [0, 1, 1] and the standard deviations are set to [1, 2, 1], respectively. For the uniform distribution, the range for each dimension is defined with lower bounds [0, -1, 1] and upper bounds [1, 1, 1.5]. The experiment results in Figure 15 show that although the initial states have high variability, the estimated values of the latent factors corresponding to the 200 individuals are ultimately highly classified and can be divided into 4 groups.Ablation Study: Consideration of transformer as encoder Our framework is flexible enough to integrate various encoders and decoders, depending on the application tasks. To demonstrate this flexibility, we incorporated Transformers into our framework and conducted a comparative analysis against the existing models. The result, shown in Figure 15(c) , indicates that while both frameworks achieve identifiability, the Transformer-based encoder demonstrates faster convergence compared to our previous approach.Added Experiment: Inventory Inventory management [71] is an important real-world problem that aims to keep inventories of goods at optimal levels to minimize inventory costs while maximizing revenue from demand fulfillment. We tested the performance of our algorithm on the inventory with state dimensions of 50, 100, and 200 and added additional baselines (8) Meta gradient RL, (9) Multitask RL, (10) Policy distillation, and (11) Non-policy adaptation to verify the model. The experimental results in Figure 12 show that our framework outperforms other algorithms in terms of initial reward and final reward.Added Experiment: AhnChemo AhnChemoEnv [61] is designed to simulate cancer treatment through chemotherapy, allowing realistic modeling of tumor growth and response to treatment. We create different groups with PK/PD variation. The experimental results in Figure 13 show that our framework outperforms other algorithms in terms of initial and final reward. Our method achieves the highest initial and final rewards compared to the baselines. Specifically, it shows a significant jump-start compared to non-policy adaptation, validating the effectiveness of our adaptation approach.The meta-gradient method optimizes the hyperparameters of the learning algorithm by calculating the gradient of the learning process, allowing rapid adaptation to new tasks as they change. However, due to the continuous adjustment of learning strategies during training, it converges more slowly and the adaptation effect is less significant compared to our algorithm. Multitask RL improves learning efficiency by sharing model strategies across different tasks. This requires first training policies on multiple tasks, which can be time-consuming (and even risky) during exploration. Moreover, identifying which new task corresponds to a previously trained task can be challenging. Our algorithm addresses this by estimating directly without requiring prior knowledge. Policy distillation transfers the knowledge of already trained teacher models to a student model, allowing the student to perform well across multiple tasks. However, this approach highly relies on the performance of the teacher models; insufficiently trained teacher models can negatively impact the final performance. Our algorithm does not depend on the source policy performance; subsequent policy optimization is based on the new environment, leading to better final performance.The estimation framework is trained using AdamW optimizer for a maximum of 200 epochs and early stops if the validation ELBO loss does not decrease for ten epochs. A learning rate of 0.001 and a mini-batch size of 32 are used. We used three random seeds in each experiment and reported the mean performance with standard deviation averaged across random seeds. We used a machine with the following CPU specifications: 11th Gen Intel(R) Core(TM) i7-11800H @ 2.30GHz with 16 logical processors. The machine has one GeForce RTX 3080 GPU with 32GB GPU memory. Our work has a significant impact on ethics, society, and future applications. We emphasize the importance of individualized policies in systems and advocate for a deeper understanding and respect for individual differences. Tailoring interventions to different individuals has the potential to improve user experience and outcomes in healthcare, education, and other areas. This approach avoids a one-size-fits-all policies. Our method can greatly improve individualized services, transforming the delivery of educational content, the management of healthcare, and the recommendation of products. This makes these services more effective and aligned with individual needs. However, the implementation of this method requires careful consideration of privacy and data security, as personalized systems require the collection and analysis of personal data. Maintaining user trust, preventing misuse, and ensuring ethical use of such data are of utmost importance. Figure 13 : Results in AhnChemoEnv. We evaluated the performance of our method against several baselines, including (1) meta gradient RL, (2) multitask RL, (3) policy distillation, and (4) non-policy adaptation. Our method outperforms these benchmarks and achieves superior performance in terms of initial reward and final reward. The proposed framework is customized based on the requirements of the identifiability theorems given in Section 4. We would like to emphasize that our proposed framework differs from the traditional VAE and model-based RL in three main aspects: (1) Our framework uses a quantization layer to discretize the continuous latent representations. This mapping of continuous latent representations to an embedding dictionary is well suited to the group determinacy requirement. (2) Our decoder reconstructs individualized state transition processes to simulate the data generation process, incorporating additional conditions as well as the estimated latent factor. (3) We further extract latent factors for each individual as additional information to facilitate individual policy learning. The detailed implementations of each component are summarized below.Encoder For any individual m, the Conv1D layer transforms an input sequence s m t , using learned kernel filters. These filters slide over the sequence to produce a feature map, denoting the response of the filter at each position. Mathematically, the transformation by a single filter in the Conv1D layer at time t is described as o t = σ W * s m t:H+t + b , where o t is the feature map, W the kernel to be learned during training, * the convolution operation, s m t:H+t the input sub-sequence from time t to t + H, where H is the size of the kernel. σ is the activation function, and b is the bias term to be learned during training. The layer may contain multiple such filters, each learning different features of the input sequence. The resulting feature maps serve as a transformed representation z m , which embeds the information about the latent group factor κ. As for the LSTM, let the hidden states and cell states of the LSTM at time t denote as h t and c t , respectively. Then, the LSTM updates are given by:f t = σ(W f • [s m t , h t-1 ] + b f ), i t = σ(W i • [s m t , h t-1 ] + b i ), ct = tanh(W c • [s m t , h t-1 ] + b c ), c t = f t ⊙ c t-1 + i t ⊙ ct , o t = σ(W o • [s m t , h t-1 ] + b o ), h t = o t ⊙ tanh(c t ),where σ is the sigmoid activation function, ⊙ element-wise multiplication. W f , W i , W c , W o and b f , b i , b c , b o are the weight matrices and bias terms to be learned during training. f t , i t , ct , c t , o t and h t are the forget gate, input gate, candidate cell state, cell state, output gate, and hidden state at time t, respectively. The final hidden state of LSTM h T , after the sequential processing of the entire trajectory, serves as the representative z m that embeds the information about the latent group factor κ.Quantization Layer Let the output of the encoder be a continuous latent representation denoted as z m ∈ R, and define an embedding dictionary E consisting of G vectors, where each vector represents a unique discrete category: E = {e 1 , e 2 , . . . , e G }, where e i ∈ R. The quantized vector κm is obtained by mapping z m to the nearest dictionary vector. The mapping can be expressed mathematically as κm = arg min ei∈E ∥z m -e i ∥ 2 . Subsequently, the quantized output is the vector from the dictionary that is closest to the encoder output. Thus, the continuous representation z m is effectively mapped to a discrete κm by finding the nearest neighbor in the dictionary, aligning the representation learning with the discrete nature of the latent variable.Decoder Suppose s m t-1 and a m t-1 as the true previous state and action, respectively. Let κm be the approximated latent group factor for the m-th individual. The inputs to the conditional decoder are a combination of the aforementioned variables: Input t = (s t-1 , a t-1 , κm ). The output of the decoder is the reconstructed next state, ŝt , which is a function of the decoder input: ŝt = De(Input t ). The reconstruction likelihood measures how closely the reconstructed state matches the true subsequent state, which is defined as L Recon = p Recon (s m t |s m t-1 , a m t-1 , κm ). The objective in this process is to optimize the decoder parameters to maximize the reconstruction likelihood max L Recon so that the reconstructed state ŝt is as close as possible to the true next state s t .The pseudocode for the proposed algorithm is presented in Algorithm 1 and Algorithm 2. Estimate individual-specific factor: κm ← Quantization(z m ) end for 13: until convergence generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 


Decision: This paper demonstrates the identifiability of latent factors and introduces a practical method for learning these processes from observed state-action trajectories. Experiments across a few [somehow similar] domains confirm that the method can effectively identify latent state-transition processes and help developing individualized RL policies. While in general the paper's presentation is good, it heavily relies on theoretical results which are provided in the appendix. As such supplementary material is not reviewed here, I would strongly recommend that the authors double-check all the results, proofs, and examples in the appendix before camera-ready version.

Reviews: review: paper_summary: paper_summary: The paper titled "Identifying Latent State-Transition Processes for Individualized Reinforcement Learning" addresses the challenge of optimizing individualized reinforcement learning (RL) policies by focusing on latent individual-specific factors that influence state transitions. This is particularly significant in domains like healthcare and education, where individual differences can causally affect outcomes.

The primary contribution of the paper is the establishment of the identifiability of latent factors that drive individualized state-transition processes. The authors propose a novel framework called Individualized Markov Decision Processes (iMDPs) that incorporates these latent factors into the RL framework. The key advancements include:

1. The paper provides theoretical guarantees for the identifiability of latent factors under both finite and infinite conditions, making it possible to distinguish different underlying components in state transitions.

2. The authors develop a generative-based method to effectively estimate these latent individual-specific factors from observed state-action trajectories. This method employs a variational autoencoder with a vector quantization layer to discretize the latent space and ensure accurate factor estimation.

3. The proposed method is empirically validated across various datasets, demonstrating its effectiveness in identifying latent state-transition processes and improving the learning of individualized policies. 

4. The paper outlines a two-stage approach for policy learning, where the estimated latent factors are used to tailor RL policies for individuals, thereby enhancing the policy’s adaptability to different environments and individuals.
strength_weakness: strength_weakness: **Assumptions on Latent Factors**
The paper assumes that the latent individual-specific factors are time-invariant, which might not hold in all real-world scenarios. For instance, in healthcare, a patient’s condition can change over time, affecting the state-transition processes. Addressing this limitation could involve extending the framework to account for time-varying latent factors or providing a discussion on how to handle such scenarios.

**Limited Real-World Validation**
While the empirical validation includes synthetic datasets and a real-world dataset (Persuasion For Good corpus), the application domains are somewhat limited. The healthcare and education examples are discussed theoretically but not empirically validated with real-world data. To strengthen the paper, the authors could include experiments using real-world healthcare or education datasets to demonstrate the practical utility and robustness of their method in these critical applications. A possible healthcare platform is [1], where there are 4 environments with customizable PK/PD variation setups.

**Policy Adaptation for New Individuals**
The approach for policy adaptation for new individuals involves initializing the policy based on the estimated group factor and fine-tuning it with new interactions. However, the paper lacks detailed evaluation metrics or benchmarks to assess the effectiveness and efficiency of this adaptation process. Including more quantitative analysis and comparisons with existing methods for policy, adaptation would enhance the understanding of the benefits and limitations of the proposed approach.

### Actionable Suggestions
1. add a discussion of time-varying latent factors in your limitation section
2. Incorporate real-world experiments in healthcare or education to validate the practical utility of the method.
3. Provide detailed evaluation metrics and benchmarks for the policy adaptation process.


References
[1] Luo, Zhiyao, et al. "DTR-Bench: An in silico Environment and Benchmark Platform for Reinforcement Learning Based Dynamic Treatment Regime." arXiv preprint arXiv:2405.18610 (2024).
questions: questions: How effective is the policy adaptation process for new individuals compared to existing methods? Can the authors provide quantitative benchmarks or metrics to evaluate this process?

How does the variability in initial states across individuals impact the performance and identifiability of the latent factors? Are there specific strategies to handle high variability in initial states?
limitations: limitations: Privacy Concerns: The individualized nature of the proposed method implies that it relies on detailed personal data. This raises potential privacy concerns, especially in sensitive domains like healthcare and education. The authors should discuss how they plan to handle privacy issues, possibly suggesting the use of privacy-preserving techniques such as differential privacy.

Robustness and Reliability: Another potential impact is related to the robustness and reliability of the individualized policies. In critical applications, such as healthcare, incorrect policy recommendations could have serious consequences. The authors should address its limitations from this angle.
paper_summary: paper_summary: The authors of this paper establish the identifiability of latent state-transition processes in reinforcement learning (RL) and propose a practical method for learning these processes from observed state-action trajectories. The focus is on personalized reinforcement learning (RL), where different individuals may exhibit different state-transition processes influenced by latent individual-specific factors. The study introduces the concept of a personalized Markov decision process (iMDP) and provides theoretical guarantees for learning state-transition processes with latent factors. Experiments on multiple datasets demonstrate the effectiveness of the method in inferring these factors and learning personalized strategies.
strength_weakness: strength_weakness: 1. **Broader Experimental Scope**: Although the experiments cover multiple datasets, including more diverse environments to test the generality of the findings would be valuable, such as testing in more complex high-dimensional RL environments.
2. **Comparison with Other Methods**: A more comprehensive comparison with other state-of-the-art methods in meta-RL and multi-task reinforcement learning could strengthen the validation of the proposed method.
3. **Discussion in Appendix**: Adding a discussion in the appendix about the differences and connections between iMDP and meta-RL, context-based RL, and multi-task reinforcement learning would be beneficial.
questions: questions: 1. **Consideration of Transformer as Encoder**: In Section 5.1, "temporal dependency from the sequential observations," the use of Transformers as the primary encoder was not considered. Supplementing with relevant experimental comparisons would provide clearer insights for readers.
limitations: limitations: The paper explicitly discusses its limitations.
paper_summary: paper_summary: The authors consider a problem where the underlying dynamical process is an MDP but with a time-invariant latent factor. The authors provide examples of such problems in the real-world. To solve these problems, the authors propose a new mathematical framework called Individualized Markov Decision Processes (iMDPs). Furthermore, the authors provide theoretical guarantees and an algorithm for learning these latent factors from observational data. The authors then demonstrate the performance of this approach on multiple datasets.
strength_weakness: strength_weakness: The authors could have considered simultaneous estimation and control, which is often needed in real-world problems. Furthermore, in choosing baselines for comparisons, the authors have not considered RL methods that take into account trajectory history.
questions: questions: 1. How does this compare with meta learning or transfer learning or a POMDP where each individual is a "slightly different" environment? While the authors state that these models are different, in terms of learning and practical considerations, how do these frameworks differ?
2. Does identifiability of these latent factors get influenced by the policy being followed to generate the data?
3. Why is the individual-specific latent factor $\kappa$ not included in the reward function?
4. What does the term $\mathbb{P}(s, a | u)$ imply in equation (1)? Since $u$ is not observed by the decision maker, how can its action $a$ be conditioned on $u$?
5. $d_s$ and $d_a$ are not defined in line 110. Can the authors add these definitions?
6. Is this learning online or offline? From Line 246 it appears that the authors only consider the offline case. However, their control part requires an online setup as mentioned in Line 335. Specifically is the learning of $\kappa$ and optimal policy simultaneous?
7. If an online setup was available, how would the authors propose to use exploration to better identify $\kappa$? Will this exploration to learn $\kappa$ be harmful/costly?
8. Can the authors define the term "post-nonlinear temporal" model in Line 162?
9. Can the authors define the covariance matrices $\Sigma_{\mathbf{A_i}, \mathbf{B_i}}$ in Line 173?
10. Is $n$ in Theorem 4.3 same as $d_s$ in Line 110?
11. Is the iMDP assumed to be an infinite horizon setup?
12. Why are the quantization loss and commitment loss defined after equation (4) the same? Also in Line 9 in Algorithm 2.
13. Is the reconstruction loss differentiable? Can the authors explain as it involves a quantization step as well? Or does this term only compute the gradient with respect to parameters of the decoder?
14. Is the policy adaptation step specified equivalent to augmenting the input given to the policy function or Q function with $\kappa$?
15. In the caption for Figure 3, it would be helpful if the authors could define the shaded regions around the curves.
16. While considering RL baselines, why haven't algorithms for POMDPs that take the entire history into account been considered?
17. Can the authors explain the motivation behind the ablation study starting in Line 317?
18. What do the authors mean by the limitation about "instantaneous causal influences" in Line 389?
19. Can the authors give more details in the proof of Theorem 4.2? (Line 678 onwards)
20. In Line 678, is $n$ the number of state components?
21. I could not follow the proof for Theorem 4.3. There are several concepts introduced in the appendix and often the proof refers to a concept explained later. It would be helpful if the authors could provide some motivation/intuition for these various constructions/concepts.
22. In line 912, arg max is a set operator and hence $=$ should be replaced with $\in$.
23. In Line 7 in Algorithm 1, what is the policy used to collect the trajectory samples?
24. Can the authors give the specific values for the dimensions and other architecture details used in their numerical examples?
limitations: limitations: The authors discuss the limitations of their work in the Conclusions section and in the Impact Statement section (Appendix H).
paper_summary: paper_summary: The paper addresses the challenge of identifying individualized state-transition processes in reinforcement learning (RL) when individual-specific factors are latent. The authors propose a novel framework, the Individualized Markov Decision Processes (iMDPs), to incorporate these latent factors into the state-transition processes. They provide theoretical guarantees for the identifiability of these latent factors and present a practical method for learning them from observed state-action trajectories. Experiments on various datasets demonstrate the effectiveness of their method in identifying latent state-transition processes and optimizing individualized RL policies.
strength_weakness: strength_weakness: 1. The method might be challenging to generalize to all types of RL problems, especially those with instantaneous causal influences within states.
2. Some sections could be simplified for broader accessibility, and future work should address the limitations mentioned.
questions: questions: 1.How does the proposed framework scale with larger datasets and more complex environments? Are there any computational limitations or bottlenecks that need to be addressed?
2. Can the framework handle real-time adaptation in dynamic environments where latent factors might change over time? If so, how would this be implemented?
limitations: limitations: 1. The framework assumes that the latent individual-specific factors are time-invariant. In real-world scenarios, individual-specific factors can evolve over time, which this model does not account for.
2. The framework might face challenges in scaling to environments with very high-dimensional state and action spaces, which are common in some real-world applications.

Comments: title: Could you please let us know whether our responses properly addressed your concern?
Could you please let us know whether our responses properly addressed your concern?
Thanks to All Reviewers. Are There Any Other Concerns?
comment: Thanks for these clarifications!
Thank you so much for your time and feedback. We sincerely appreciate you updating your recommendation. Please let us answer your questions below.

**Follow-Up Q1: If the states are fully observable, then in an MDP framework, the transitions should not depend on any other factor. Am I missing something here?**

**Follow-Up A1**: Thanks for the good question. According to the definitions [1-2], the difference between POMDPs and MDPs is that in POMDP, the agent is unable to directly observe the current state $s$. Instead, the POMDP agent only receives a noisy or partial observation $o$, which is decided by the sensor model $O(o|s, a)$ (or sometimes $O(o|s)$). 

In our problem setting, although the states $s$ are fully observable, the transition probability $\mathbb{P}(s’|s, a, \kappa)$ is influenced by unobserved variable $\kappa$, along with the observed variables ($s$ and $a$). The dependency of the transition on observed or unobserved variables does not necessarily determine whether the environment is modeled as an MDP or a POMDP. Therefore, despite including a latent variable, our setting conforms to the definition of an MDP.

[1] Kaelbling, Leslie Pack, Michael L. Littman, and Anthony R. Cassandra. "Planning and acting in partially observable stochastic domains." Artificial intelligence 101.1-2 (1998): 99-134.

[2] Igl, Maximilian, et al. "Deep variational reinforcement learning for POMDPs." International conference on machine learning. PMLR, 2018.


**Follow-Up Q2: Regarding A12 and A13, I find the definitions to be identical. What am I missing here?**

**Follow-Up A2**: Sorry for the confusion and you are right. Thank you for pointing out this typo, and we have corrected it in the updated manuscript. The quantization loss and commitment loss should be 

$\mathcal{L} _ {\text{Quant}} = \sum_i \left\|\| \text{sg}[z _ {m,i}] - e _ {m,i} \right\|\|^2$, and 

$\mathcal{L} _ {\text{Commit}} = \sum_i \left\|\| z _ {m,i} - \text{sg}[e _ {m,i}] \right\|\|^2$, respectively.
I thank the authors for their detailed explanations to my questions. Based on these responses, I have increased my score. However, there are some points that I have not yet understood:

A. "POMDPs focus on dealing with incomplete information and making decisions under uncertainty with time-varying hidden states. Our work emphasizes individualized decision-making considering latent factors that influence state transitions with fully observable states". I did not understand this. If the states are fully observable, then in an MDP framework, the transitions should not depend on any other factor. Am I missing something here?

B. Regarding A12 and A13, I find the definitions:
1. $\mathcal{L}_{Quant} = \sum_i \lVert sq[z_i] -  e_i \rVert^2$ 
2. $\mathcal{L}_{Commit} = \sum_i \lVert e_i - sq[z_i] \rVert^2$ 
to be identical (I have left out the $m$ in the subscript on the RHS due to a formatting error. What am I missing here?
Dear Reviewer cEn6,

Thank you so much for your time and feedback. It means a lot to us. We are so happy that most of your major concerns were properly addressed. 

Wish you all the best,

Authors of submission 12027
Thank you for the response. The authors have addressed most of my major concerns.
I have no further questions, and I am happy to raise the score to 7.
Dear Reviewer t6CB,
 
Thank you again for your valuable time dedicated to reviewing our paper and for your helpful suggestions. We particularly appreciate your questions regarding the problem setting, theorem, and framework. So we are eager to see whether our responses properly addressed your concerns and would be grateful for your feedback.
 
With best wishes,
 
Authors of submission 12027
Dear Reviewer pu6U,
 
Thank you again for your valuable comments. Your suggestions on the experiments were very helpful, and we are eager to learn whether our additional experimental results and our responses have properly addressed your concerns.
 
Due to the limited time for discussion, we hope to see your feedback and hope for the opportunity to respond to your further questions, if there are any.
 
Yours sincerely,
 
Authors of submission 12027
Dear Reviewers,

Thank you very much for your time and effort in reviewing our manuscript. We have tried our best to address all your concerns and have revised our paper accordingly. 

As the discussion deadline approaches, please let us know if you have any further concerns and we will be happy to address them.

Sincerely,

Authors of Submission #12027