Identifying Latent State-Transition Processes for Individualized Reinforcement Learning

Q1: The method might be challenging to generalize to all types of RL problems, especially those with instantaneous causal influences within states.

A1: Thank you for pointing out this scenario. It is indeed possible that there are instantaneous causal influences within states in the considered problem. However, even with these influences, our framework can still identify the latent individual-specific factors, and the theoretical results will not be affected. In the instantaneous case, since the states are observed, we can treat the states at each time step as a whole and input them into the current estimation framework. Thus the instantaneous causal influences within states will not affect our current estimation framework and theorem results.

Q2: Some sections could be simplified for broader accessibility, and future work should address the limitations mentioned.

A2: Thank you for your suggestion. In the updated manuscript, we have reorganized the sections to improve accessibility and clarity. Additionally, we have added a separate Limitations section to discuss how we address limitations, including instantaneous causal influences.

Future work could extend our current framework from a causal perspective. For instance, we could develop modules that explicitly account for causal dependencies within states, using techniques such as causal graphical models and advanced inference methods that handle instantaneous causal relationships [1]. Such extensions could be directly integrated into our current framework, and we have been working on this problem.

[1] Li, Zijian, et al. "On the Identification of Temporally Causal Representation with Instantaneous Dependence." arXiv preprint arXiv:2405.15325 (2024).

Q3: 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? 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.

A3: In the updated manuscript, we validated our algorithm in the AhnChemoEnv in DTRGym [2] and inventory management tasks [3]. AhnChemoEnv is designed to simulate cancer treatment through chemotherapy, allowing realistic modeling of tumor growth and response to treatment. Inventory management 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 test the performance of our algorithm on the inventory with state dimensions of 50, 100, and 200. The experimental results (see Figure 1 and Figure 3 in REBUTTAL.pdf) show that our framework outperforms other algorithms in terms of initial reward and final reward.

Specifically, our method shows a significant jump-start compared to non-policy adaptation, validating the effectiveness of our approach. Unlike the meta-gradient method, which requires continuous adjustment of learning strategies during training and thus converges more slowly with less significant adaptation effects, our algorithm is more efficient. Additionally, training across multiple tasks in multitask RL can be time-consuming, and identifying which new task corresponds to a previously trained task can be challenging. Our algorithm addresses this by directly estimating $\kappa$ without requiring prior knowledge. Furthermore, policy distillation heavily depends on the performance of teacher models; insufficiently trained teacher models can negatively impact the final performance. Our algorithm, however, does not rely on source policy performance. Instead, it optimizes the policy based on the new environment, leading to better final performance.

Our findings suggest that while the algorithm scales reasonably well, certain computational limitations, such as increased processing time and memory consumption, need to be addressed. In high-dimensional scenarios, the use of hardware accelerators such as GPUs, as well as fine-tuning the model, is essential to maintain optimal performance.

[2] 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).

[3] Sun, Yuewen, et al. "ACAMDA: Improving Data Efficiency in Reinforcement Learning Through Guided Counterfactual Data Augmentation." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 38. No. 14. 2024.

Q4: Can the framework handle real-time adaptation in dynamic environments where latent factors might change over time? If so, how would this be implemented? 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.

A4: Thank you for pointing out this interesting scenario. It is indeed possible for the latent individual-specific factor to be time-variant in the considered problem. We believe our framework can be extended to handle time-varying cases, although establishing the theoretical identifiability is highly non-trivial. For instance, if we directly allow these latent factors to be time-varying, it seems hopeless to recover the latent variables, since each individual may have a specific latent factor. Therefore, further constraints would be needed. For instance, the theoretical identifiability might benefit from the constraint that the total number of possible values of the latent factors over time and across different individuals is finite. We hope this important extension can be achieved by researchers in the field by leveraging our framework and further considering the idea in [4].

[4] Hu, Yingyao, and Matthew Shum. "Nonparametric identification of dynamic models with unobserved state variables." Journal of Econometrics 171.1 (2012): 32-44.

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

Q1: Broader experimental scope: Thank you very much for your suggestion. In the updated manuscript, we validated our algorithm in the AhnChemoEnv in DTRGym [1] and inventory management tasks [2]. AhnChemoEnv is designed to simulate cancer treatment through chemotherapy, allowing realistic modeling of tumor growth and response to treatment. Inventory management 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. The experimental results (see Figure 1 and Figure 3 in REBUTTAL.pdf) show that our framework outperforms other algorithms in terms of initial reward and final reward. Detailed analyses are provided in Q2.

[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).

[2] Sun, Yuewen, et al. "ACAMDA: Improving Data Efficiency in Reinforcement Learning Through Guided Counterfactual Data Augmentation." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 38. No. 14. 2024.

Q2: Comparison with other methods: We appreciate your suggestion. In the updated manuscript, we have added more benchmarks and further compared our algorithm against (1) meta gradient RL, (2) multitask RL, (3) policy distillation, and (4) non-policy adaptation. Additionally, we have included two evaluation metrics: (1) initial reward, which measures the initial performance benefited from policy adaptation, and (2) final reward, which measures the performance after the full training process. Please refer to REBUTTAL.pdf for the detailed results.

The comparisons suggest that our method achieves the highest initial and final rewards compared to the benchmarks. 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 using $\kappa$ 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.

Q3: Discussion in Appendix: In the updated manuscript, we have added a section discussing the differences and connections between iMDP and meta-RL, context-based RL, and multi-task reinforcement learning, summarized below.

• Meta-learning 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. 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.

• Contextual MDPs 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 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. It often involves sharing information between tasks to improve learning efficiency and policy performance. Instead of focusing on policy optimization for all tasks, iMDP 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.

Q4: Consideration of transformer as encoder: We appreciate your insightful suggestion. In the revision, we have incorporated transformers into our framework to perform a comparative analysis against the models currently in use, and the results are reported in Figure (c) in REBUTTAL.pdf. Although both frameworks can achieve identifiability, the Transformer encoder can achieve faster convergence compared with our previous framework.

A1: Thanks for the helpful questions. Here is a brief summary of the differences. We also empirically compared our method with meta-learning and transfer-learning techniques and reported results in REBUTTAL.pdf.

• Meta-learning trains the 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.

• Transfer learning focuses on leveraging knowledge from one domain to improve performance in a new domain while facing challenges such as negative transfer. Our method can identify the latent individual-specific factors and realize individualized decision-making with interpretability.

• 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.

A2: In the current theoretical treatment, yes, the identifiability is influenced by the policy. According to our proof technique, this assumption is needed. Thank you for pointing it out and we have made this condition explicit in the updated paper. At the same time, it might be possible to avoid this assumption although the proof seems to be nontrivial.

A3: In our setting, we assume that different groups have no different reward weightings. Thus $\kappa$ only influences the states and transition dynamics.

A4: We believe this might be a misunderstanding. $\mathbb{P}(s,a|u)$ represents the joint distribution of $(s,a)$ for each individual. Here, $u$ denotes different individuals which is explicitly observed. This notation is used for mathematical convenience.

A5, A10, A20: Both $n$ in Theorem 4.3 and $d_s$ in line 110 represent the state dimension, and $d_a$ denotes the action dimension. We didn’t find $n$ in line 678–perhaps you referred to another line number? If you mean the $n$ in lines 166, 728, or 739, it also represents the state dimension.

A6, A7: The estimation network is pre-trained offline. When a new individual arrives, we estimate $\kappa$ and adapt the policy simultaneously, and the policy is adapted through new interactions. Such exploration is necessary to better identify $\kappa$ and discover the optimal policy in RL. This has been provided explicitly in the updated paper.

A8, A9: The definition of the post-nonlinear temporal model and covariance matrices are provided in Appendix C.1 and C.4.1, respectively. In light of your question, we have also moved them to the main paper in the updated paper.

A11: Yes, we assume an infinite horizon setup in iMDP.

A12, A13: This might be a misunderstanding. The quantization loss and the commitment loss are different and serve unique purposes. The key is in the positioning of the stop gradient (sg[·]) [1], showing how these losses influence the encoder and the codebook differently. The stop gradient allows the model to handle the non-differentiable quantization step, ensuring that gradients are computed for the encoder and the rest of the model, making the reconstruction loss differentiable.

[1] Van Den Oord, Aaron, and Oriol Vinyals. "Neural discrete representation learning." Advances in neural information processing systems 30 (2017).

A14: Yes, we identify and incorporate $\kappa$ into policy function to learn individualized policy.

A15: The shaded regions represent the standard deviation. We have added an explanation in the revision.

A16: Our problem and POMDP address different settings. We assume fully observable states influenced by latent individual-specific factors, rather than partially observable states. Under the similar setting, we compared our method with meta-learning and transfer learning techniques.

A17: The ablation study analyzes the contributions of different components in the framework. The sequential encoder aligns our framework with the theorem requirements, and the conditional decoder meets our data generation process. The noise estimator is used to increase robustness.

A18: Our current framework only considers transitions from time $t$ to $t+1$. However, in some cases, there may be causal influences that happen instantaneously within the same time step, which is beyond the scope of our current discussion.

A19: Individuals sharing the same latent factor can be grouped due to similar state transitions. For a specific state-action pair $(s^*, a^*)$, the next state $s'$ should be consistent across individuals in the same group. We define $t^j$ as the time of the $j$-th occurrence of $(s^*, a^*)$ and collect states $X^j$ = {$s_{t^j+1} \mid t^j = 1, 2, \dots$}. By comparing $X^j$ across different individuals, we group those with similar transition behaviors and achieve identifiability by Lemma B.1.

A21: We appreciate your feedback–indeed, this part is dense and involves many concepts and derivations. The proof for Theorem 4.3 has two parts.

• Structure Identifiability: Lemma C.1 shows that the rank of the sub-covariance matrix of the observed variables equals the number of minimal choke sets, which represent the minimal sets that separate the two observed variable sets forming the sub-covariance matrix. The difference between the calculated rank and the expected choke factors indicates the number of latent factors.

• Parameter Identifiability: We estimate $\alpha$ and $\beta$, which describe the influence of states and actions, using a regression model. The latent coefficient $\lambda$ is identified using orthogonal techniques in factor analysis.

A22: We have revised the notation accordingly.

A23: The trajectory samples are collected using a random policy.

A24: The summary of architecture details is provided in REBUTTAL.pdf and more details have been added in the updated manuscript.

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

Q1: Discussion of time-varying latent factors: Thank you for pointing out this interesting scenario. It is indeed possible for the latent individual-specific factor to be time-variant in the considered problem. We believe our framework can be extended to handle time-varying cases, although establishing the theoretical identifiability is highly non-trivial. For instance, if we directly allow these latent factors to be time-varying, it seems hopeless to recover the latent variables, since each individual may have a specific latent factor. Therefore, further constraints would be needed. For instance, the theoretical identifiability might benefit from the constraint that the total number of possible values of the latent factors over time and across different individuals is finite. We hope this important extension can be achieved by researchers in the field by leveraging our framework and further considering the idea in [1].

[1] Hu, Yingyao, and Matthew Shum. "Nonparametric identification of dynamic models with unobserved state variables." Journal of Econometrics 171.1 (2012): 32-44.

Q2: Additional real-world experiments in healthcare: Thank you very much for your suggestion and for providing the healthcare platform. In the updated manuscript, we validated our algorithm in the AhnChemoEnv environment in DTRGym [2] and created different groups with PK/PD variation. The experimental results (see Figure 1 in REBUTTAL.pdf) show that our framework outperforms other algorithms in terms of initial & final reward. Detailed analyses are provided in Q3.

[2] 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).

Q3: Evaluation and discussion on policy adaptation process: We appreciate your suggestion. In the updated manuscript, we have added more baselines and further compared our algorithm against (1) meta gradient RL, (2) multitask RL, (3) policy distillation, and (4) non-policy adaptation. Additionally, we have included two evaluation metrics: (1) initial reward, which measures the initial performance benefited from policy adaptation, and (2) final reward, which measures the performance after the full training process. Please refer to REBUTTAL.pdf for the detailed results.

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 using $\kappa$ 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.

Q5: Discussion on variability in initial states: This is a great question! We can still establish the identifiability of the latent variable for each trajectory provided by each user. However, since each trajectory has a finite (usually pretty small) length, the estimated values might be noisy. The second step is to merge latent variables that should have the same value. A straightforward method is as follows: we first estimate the value of the latent variable for each individual and then see how they can be clustered together. We check whether they should be classified into the same or different groups by conducting clustering.

A more principled way to address this issue has yet to be developed, and we have conducted additional experiments to verify the feasibility and report the results in Figure (a-b) in REBUTTAL.pdf. 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]. As you can see from the experiment, although the initial state has a high variability, finally the estimated values of the latent factors corresponding to the 200 individuals are highly classified and can be divided into 4 groups.

Q6: Impact on privacy issues & robustness and reliability: In the updated manuscript, we have included a separate limitation section to discuss how we address these limitations.

• Potential privacy risks: We could use de-identifying techniques and remove direct identifiers such as names and zip codes, apply masking techniques such as data perturbation, and use pseudonymization to replace private identifiers with artificial ones. Incorporating differential privacy techniques is also interesting and we will leave it as future work.

• Robustness and reliability: We could provide an assumption checklist to help users determine whether our work is applicable to their specific scenarios, thereby avoiding misuse and improving reliability.