IOSTOM: Offline Imitation Learning from Observations via State Transition Occupancy Matching

We thank reviewer for your helpful comments and feedback. We provide our responses below to address your concerns.

1. What is the main reason for considering mixed state visitation distribution instead of starting from directly matching the objective with $d^E$? (What is the main problem if we derive it in the same way as direct matching for $d^E$?)...

We thank the reviewer for the insightful question. The objective of directly matching between $d^{\pi}$ and $d^{E}$ requires on-policy samples from $\pi$ to solve which is not ideal in offline setting because these samples can be out-of-distribution. This issue is already discussed in previous LfO works like LobsDICE (Section 2.2 in [1]) and SMODICE (Section 3.1 in [2]). To mitigate this, these methods need to integrate the action-labeled suboptimal dataset into learning process by solving surgroate objectives as we note in Section 3 of our paper. Through our formulation in Eq. 2, we are also able to leverage the suboptimal dataset in training process. However, the key distinction is that our objective is mathematically constructed to have the same optimal solution as the original, direct-matching problem.

2. How was the value of hyperparameter $\alpha$ selected? How sensitive is IOSTOM's performance to hyperparameter $\alpha$?

Please refer to Point 5 of the responses to reviewer 4HqZ

3. The optimal value of the hyperparameter alpha seems to be different for each environment. Is there no way to automatically determine $\alpha$(also for $\beta$)?

We thank the reviewer for the questions. As pointed out in our response to Reviewer 4HqZ, $\alpha$ is not a sensitive parameter and can be fixed to 0.5 following previous works. The hyperparameter $\beta$, which controls the regularization strength in the extreme value learning framework, is often sensitive and may need to be tuned per environment to achieve the best performance, as shown in [3]. However, we found very minimal hyperparameter tuning was required in our framework. As shown in Appendix C.3, we fix $\beta$ to 20 for all locomotion tasks and $\beta$ to 10 for most manipulation tasks except for "pen-cloned" where β is set to 3. The ablation study of $\beta$ in Appendix D.3 also shows that in IOSTOM, $\beta$ is only sensitive per class of environment, not per individual environment. Specifically, high values like 15 or 20 are consistently effective for locomotion tasks, while lower values like 5 or 10 achieve good results on manipulation tasks. This provides a straightforward heuristic for setting this parameter: one can either select a value based on the task type (locomotion vs. manipulation) or perform a quick search on a single representative task from a class and apply that value to all others in the same class.

4. LobsDICE is also a representative comparison algorithm that can be considered in the same setting. Why didn't you compare it?...

We thank the reviewer for bringing up LobsDICE. Initially, we did not include it in our baselines because PW-DICE—a more recent and related method—was already shown in [4] to outperform LobsDICE on the “random+expert” D4RL setting.

To address your question and provide a more complete comparison, we have conducted additional experiments comparing LobsDICE, PW-DICE, and our method (IOSTOM), as shown in the table below.

The results clearly demonstrate that IOSTOM significantly outperforms LobsDICE across all tasks. Like other DICE-based methods, LobsDICE performs poorly on the “few-expert” and manipulation tasks, while our method remains robust. Our findings also reaffirm that PW-DICE surpasses LobsDICE in the “random+expert” scenario. However, PW-DICE underperforms in several other settings, highlighting the general instability of discriminator-based methods. These comparisons further emphasize the advantages of our stable, discriminator-free approach. We will include the full results in the updated paper.

5. In the paper, it was explained that the absence of discriminator learning is an advantage, but instead $V$, learning was introduced. Isn't it similar? What makes it better?

We thank the reviewer for the insightful question. In our setting, discriminator learning and value function learning serve fundamentally different purposes. The discriminator acts as a reward function, assigning high scores to expert states and low scores to non-expert ones. In contrast, our value function does more than binary classification—it imposes a structured ranking over suboptimal states by assigning higher values to those from which expert states are more easily reachable. This provides a meaningful recovery signal, guiding the agent back toward expert-like behavior when it deviates. Such nuanced guidance is something a simple reward-like discriminator cannot provide.

6. In all experiments, were you only training using the offline dataset? Was there no access to the test environment even during hyperparameter search and final model selection?

We thank the reviewer for the question. In our work, the entire training process for IOSTOM is conducted strictly offline, using only the provided datasets. This stands in contrast to online IL methods like GAIL [6], which require continuous interaction with the environment to populate a replay buffer.

However, we do require limited access to the environment for hyperparameter tuning and performance evaluation. This is a standard practice in both offline RL (e.g., Table 4 in [3]) and offline IL (e.g., Table 7 in [5]), where returns from the environment are used to select the best-performing hyperparameters for each task or task class.

While fully offline alternatives exist, they are less applicable in our setting. For example, off-policy evaluation (OPE) methods [7] are commonly used in offline RL but require reward information, which is not available in our imitation learning setup. Another alternative—using validation loss to guide selection—is often unreliable due to its poor correlation with final policy performance [8].

Therefore, we adopt the commonly accepted assumption of having limited evaluation access to the environment. In our case, since $\alpha$ is fixed at 0.5, we only tune $\beta$ from a small candidate set {3, 5, 7, 15, 20}, requiring at most five evaluation trials.

In terms of your concern about fair comparisions where we might change the $\alpha$ and $\beta$ per environment to beat other approaches, we want to note that we keep the same $\beta=20$ across all locomotion tasks. We only use different values of $\beta$ ($\beta=3$ for "pen-cloned" task and $\beta=10$ for the rest) when solving manipulation tasks. When dealing with different type of tasks like manipulation, previous LfO methods like SMODICE or PW-DICE also needs to modify their agorithmic choices: changing from KL divergence to $\chi^2$ divergence in SMODICE (Appendix E.2 in [2]) and change in hyperparameter in PW-DICE (Table 1 in [4]). If we set a consistent $\beta=10$ for all manipulation tasks, IOSTOM is still the leading approach 7/8 tasks. Therefore, we believe that IOSTOM outperforms all baselines under a fair and standard comparison protocol.

References:

[1]: Kim, Geon-Hyeong, et al. "Lobsdice: Offline learning from observation via stationary distribution correction estimation." NeurIPS 2022.

[2]: Ma, Yecheng, et al. "Versatile offline imitation from observations and examples via regularized state-occupancy matching." ICML 2022.

[3]: Garg, Divyansh, et al. "Extreme q-learning: Maxent rl without entropy." ICLR 2023.

[4]: Yan, Kai, Alexander G. Schwing, and Yu-Xiong Wang. "Offline imitation from observation via primal wasserstein state occupancy matching." ICML 2024.

[5]: Chi, Cheng, et al. "Diffusion policy: Visuomotor policy learning via action diffusion." The International Journal of Robotics Research (2023)

[6]: Ho et al. "Generative adversarial imitation learning." NeurIPS 2016.

[7]: Fu, Justin, et al. "Benchmarks for deep off-policy evaluation." ICLR 2021.

[8]: Mandlekar, Ajay, et al. "What matters in learning from offline human demonstrations for robot manipulation." CoRL 2021.

We hope our revisions and our response  address the reviewers’ concerns and further clarify our contributions. If there are any additional questions or comments, we would be happy to address them.

Thank you for the detailed response from the authors. Most of my questions were answered through the author's response. However, I do not completely agree with the claim that it is not sensitive to the hyperparameter alpha, and I still have concerns that there is no guarantee that it will work regardless of the value of alpha when experimenting in a new domain. I will also check all the other reviews and responses and decide whether to change the ratings or not.

We thank the reviewer for the constructive comments and valuable feedback. Below, we provide our detailed responses to address your concerns.

1. The training objective in Eq. 9 seems potentially unstable, as there is a term $(Q(s, s')-\frac{2}{1-\gamma})^2$. When $\gamma$ is close to 1 with long horizon tasks, this can lead to a very large value of $Q$ which causes numerical issues during optimization.

We thank the reviewer for the insightful comment. In our work, the discount factor $\gamma$ is set to 0.99 (see Table 5 in the Appendix), which is a standard choice in RL. With $\gamma = 0.99$, the term $(Q(s, s') - \frac{2}{1 - \gamma})^2$ becomes $(Q(s, s') - 200)^2$. Our training remains stable with this formulation. Prior imitation learning methods such as LS-IQ [1] have also adopted this formulation to promote stable learning. The goal of $(Q(s, s') - \frac{2}{1 - \gamma})^2$ is only assigning large and bounded $Q$ values to expert transitions as pointed in Section 5 and Appendix B.8 of our work. Therefore, one could simply cap the target value (e.g. at 200) in the case of $\gamma$ > 0.99 (which is uncommon in deep RL). This would preserve the objective's intent to assign a high and stable Q-value to expert data while avoiding any potential numerical issues.

2. The number of expert dataset used in the experiment is quite large, especially for the MuJoCo datasets. In the appendix, the authors stated that they use 30 expert trajectories for "few-expert" scenario. In contrast, both SMODICE (https://arxiv.org/pdf/2202.02433, Sec. 6.1) and PW-DICE (https://arxiv.org/pdf/2311.01331, Appendix D.2) uses single expert trajectory. I would like to see how the proposed algorithm work with even less amount of expert trajectories than the current setting.

We thank the reviewer for raising this point. We believe there may be a misunderstanding. In Appendix C.1, we stated: “The state-only expert dataset in all tasks includes only one expert trajectory.” In our paper, the term “few-expert” refers to a suboptimal dataset constructed by combining "random" or "medium" trajectories with 30 expert trajectories. This is similar to the experimental setting used in SMODICE and PW-DICE, where the suboptimal dataset contains 200 expert trajectories alongside lower-quality ones. The key difference across scenarios lies in how the suboptimal dataset is constructed, as detailed in Table 4 of the Appendix.

We would greatly appreciate it if the reviewer could point us to the specific part that caused confusion, so we can revise it for greater clarity.

3. Is it possible to add a teaser figure for this paper? I feel people who are less familiar with the state occupancy matching framework can understand the idea more easily with an illustration.

We thank the reviewer for the valuable suggestion, which could indeed improve the readability of the paper. We will include such a teaser figure in the revised version. Due to the NeurIPS rebuttal policy, we are unable to upload images at this stage, but it will be added in the updated version of the paper.

4. Is there a more direct ablation / empirical evidence (e.g. visualization) to show in some cases discriminators fail other than final performance difference?

We thank the reviewer for the valuable suggestion. Visualizing the reward landscape is indeed a helpful way to illustrate how the discriminator assigns high rewards to non-expert states in tasks where discriminator-based methods perform poorly. As the NeurIPS review system does not allow image uploads during the rebuttal phase, we will include these visualizations in the revised version of the paper.

5. Why do the authors choose $\alpha =0.5$ for distribution mixture with suboptimal dataset, which seems quite large? Can the authors provide an ablation on different values of $\alpha$?

We thank the reviewer for the insightful question. We set $\alpha = 0.5$ as it is a commonly adopted choice in prior works [1, 2, 3] that employ mixture distributions for $f$-divergence minimization.

In response to your question, we conducted an ablation study on the sensitivity of $\alpha$, testing five values across all 24 tasks. The table below reports performance for each setting, with the best-performing $\alpha$ for each task highlighted in bold. We also include two additional metrics: $gap_{worst}$, representing the percentage gap between the best and worst $\alpha$, and $gap_{default}$, indicating the gap between the best-performing $\alpha$ and our default setting of $\alpha = 0.5$.

According to the table, the average $gap_{worst}$ value across all tasks is just 4.54%, which is relatively small. This indicates that IOSTOM's performance is not highly sensitive to the choice of $\alpha$ . Furthermore, the average performance gap between the task-specific optimal $\alpha$ and our default choice of $\alpha=0.5$ is even smaller at 2.07%. The $gap_{default}$ value is also under 5% in all but 3 of 24 tasks. This confirms that $\alpha=0.5$ is a robust and effective hyperparameter choice, consistently providing near-optimal performance.

6. Why do the authors choose to update policy during the training loop of other components (as shown in Alg. 1) instead of finishing the other parts of training first, which could potentially cause instability? $\theta$ is not involved in other parts of training after all.

We thank the reviewer for the insightful question. The reason is that the policy does not influence the value function updates, so policy extraction can be performed either concurrently with or after value learning. In our implementation, we follow the XQL [4] codebase, which updates the policy jointly with the value networks. This joint training approach is also adopted by other LfO methods such as DILO and LobsDICE (see Section 4.3 of [3] and Appendix G.2 of [5]).

References:

[1]: Al-Hafez, Firas, et al. "Ls-iq: Implicit reward regularization for inverse reinforcement learning." ICLR 2023.

[2]: Sikchi, Harshit, et al. "Dual rl: Unification and new methods for reinforcement and imitation learning." ICLR 2024.

[3]: Sikchi, Harshit, et al. "A dual approach to imitation learning from observations with offline datasets." CoRL 2024.

[4]: Garg, Divyansh, et al. "Extreme q-learning: Maxent rl without entropy." ICLR 2023.

[5]: Kim, Geon-Hyeong, et al. "Lobsdice: Offline learning from observation via stationary distribution correction estimation." NeurIPS 2022.

We hope our revisions and our response  address the reviewers’ concerns and further clarify our contributions. If there are any additional questions or comments, we would be happy to address them.

We thank reviewer for your helpful comments and feedback. We provide our responses below to address your concerns.

1. Access to fully observable actions in sub-optimal dataset is hard to guarantee. It would be helpful to discuss how realistic it is.

We thank the reviewer for the insightful comments. We believe that access to action-labeled suboptimal datasets is highly realistic in many practical settings. For example, in the robotics community, collecting demonstrations via teleoperation from non-expert users is often significantly easier and more cost-effective than acquiring expert demonstrations, which typically require domain-specific training or expert operators [1].

Beyond robotics, similar trends are observed in other domains. In autonomous driving, for instance, large-scale datasets are often gathered from regular human drivers who may not always exhibit optimal behavior. These action-labeled suboptimal logs still provide valuable data for training and evaluation. In healthcare applications, clinician interactions recorded in electronic health records can be viewed as suboptimal demonstrations, yet they form the basis for many decision-support models. Similarly, in gaming and digital environments, user gameplay logs—despite being noisy or suboptimal—are widely used for imitation and reinforcement learning research.

These examples underscore that suboptimal, action-labeled data are not only easier to obtain but also more abundant and representative of real-world behavior. As such, they offer a practical and scalable foundation for learning from demonstrations in many domains.

We will incorporate these examples into the revised paper to better illustrate the practicality of accessing action-labeled suboptimal datasets.

2. What is the intuition of using a mixture of joint state-visitation distributions $d_{mix}^I, d_{mix}^{E,I}$? It would be helpful to introduce the intuition before deriving the full equation.

We thank the reviewer for the insightful question. The use of a mixture of joint state-visitation distributions is adopted from some recent state-of-the-art papers on RL and IL (e.g., DualRL [2]), which highlights its advantage in avoiding on-policy sampling.

Specifically, prior work typically adopts matching between $d^{\pi}$ and $d^{E}$, but this approach requires on-policy samples from $\pi$ to solve. This is not ideal in the offline setting, as these samples can be out-of-distribution. This limitation has been discussed in previous LfO works such as LobsDICE (Section 2.2 in [3]) and SMODICE (Section 3.1 in [4]). Using a mixture of joint state-visitation distributions helps address this issue by avoiding reliance on on-policy data while maintaining effective alignment between expert and policy behavior. We will include this discussion earlier in the updated paper to better clarify the motivation behind this design choice.

3. How does the mixing parameter $\alpha$ influence the training?

We thank the reviewer for the question. In response, we conducted an ablation study on the sensitivity of $\alpha$, testing five values across all 24 tasks. The table below reports performance for each setting, with the best-performing $\alpha$ for each task highlighted in bold. We also include two additional metrics: $gap_{worst}$, representing the percentage gap between the best and worst $\alpha$, and $gap_{default}$, indicating the gap between the best-performing $\alpha$ and our default setting of $\alpha = 0.5$.

According to the table, the average $gap_{worst}$ value across all tasks is just 4.54%, which is relatively small. This indicates that IOSTOM's performance is not highly sensitive to the choice of $\alpha$. Furthermore, the average performance gap between the task-specific optimal $\alpha$ and our default choice of $\alpha=0.5$ is even smaller at 2.07%. The $gap_{default}$ value is also under 5% in all but 3 of 24 tasks. This confirms that $\alpha=0.5$ is a robust and effective hyperparameter choice, consistently providing near-optimal performance.

References:

[1]: Lynch, Corey, et al. "Learning latent plans from play." Conference on robot learning. Pmlr, 2020.

[2]: Sikchi, Harshit, et al. "Dual rl: Unification and new methods for reinforcement and imitation learning." ICLR 2024.

[3]: Kim, Geon-Hyeong, et al. "Lobsdice: Offline learning from observation via stationary distribution correction estimation." NeurIPS 2022.

[4]: Ma, Yecheng, et al. "Versatile offline imitation from observations and examples via regularized state-occupancy matching." ICML 2022.

We hope our revisions and our response  address the reviewers’ concerns and further clarify our contributions. If there are any additional questions or comments, we would be happy to address them.

Dear Reviewer gGMk

We sincerely thank the reviewer for carefully reading our paper, providing a thoughtful response, and offering valuable feedback. We will certainly incorporate the corresponding discussions and the additional experiments. We remain fully available and would be glad to provide any further details or clarify any remaining points.

We believe our responses have addressed most of the concerns raised by you and the other reviewers. If our replies and the planned revisions have satisfactorily resolved your concerns, we would be grateful if you could kindly reflect this in your evaluation of our work.

Thank you once again for your constructive engagement.

We thank reviewer for your helpful comments and feedback. We provide our responses below to address your concerns.

1. Issues about paper's clarity

We thank the reviewer for the valuable comments regarding the clarity of our paper. We will revise the manuscript accordingly based on your suggestions.

In Section 3, the definitions of $d^E(s), d^I(s)$ and are missing or unclear.

They are occupancy measures of the expert and sup-optimal datasets. We will clarify them.

In Section 4.1, $d_0(s)$ is not properly defined.

This is the visitation distribution of state $s$ at time 0. We wil clarify this.

The transition from Equation (3) to (4) is not clearly explained. I was able to derive (4) from (3), but it required effort that could have been avoided with a brief clarification.

We have removed some intermediate derivations due to space limitations. We will clarify this transition in the revised version.

The paper does not specify where the proofs of the main theorem can be found, even though they are in the appendix.

Thank you for the comment. We will add references to the corresponding proofs in the appendix.

Additionally, several parts of the text (e.g., lines 121, 124, 146, 148, 150, 225, 235, 244) are written in a way that makes them difficult to parse.

Thank you for pointing out these detailed issues. These parts were compressed due to space limitations, which may have caused some lack of clarity. We will revise them to make the content clearer and easier to follow.

2. Have the authors conducted any experiments with subsampled random data to assess the sensitivity of the method to the amount of random data?

We thank the reviewer for the insightful question. In response, we applied subsampling rates of 5 and 20, meaning only 20% and 5% of the low-quality trajectories ("random" or "medium") from the original suboptimal dataset were used, respectively. The table below compares the performance of IOSTOM(full) with these subsampled variants—IOSTOM(sub5) and IOSTOM(sub20). Drops in performance exceeding 5% relative to the full dataset are highlighted in bold.

When reducing the low-quality data to 20% of its original size (IOSTOM(sub5)), our method demonstrates strong robustness. Across all tasks and datasets, the performance remains very close to the results on the full dataset across all tasks with no significant degradation observed. When the low-quality data is aggressively subsampled to just 5% (IOSTOM(sub20)), we observe a more noticeable performance drop (up to about 18%), particularly in th halfcheetah and ant environments. However, IOSTOM is still robust on 10/16 tasks in this extremely limited data scenario.

3. Have the authors conducted any experiments in image-based environments? Such settings may better reflect the nature of imitation from observation scenarios, where actions are not available and only visual state information is provided.

We thank the reviewer for the insightful question. In this work, we focused on the well-established D4RL state-based benchmarks to allow for direct and rigorous comparison with recent state-of-the-art LfO methods (e.g., SMODICE, PW-DICE, DILO), which are also primarily evaluated in these settings.

While we agree that extending to image-based tasks would broaden the applicability of our method, experiments on domains like V-D4RL [1] require careful design of visual encoders and significant effort to ensure fair integration with both our framework and existing baselines. As most prior LfO methods also do not evaluate in image-based settings, a thorough and fair comparison remains challenging.

We will include a discussion in the revised paper highlighting the extension to vision-based tasks as a key direction for future work.

References:

[1]: Lu, Cong, et al. "Challenges and opportunities in offline reinforcement learning from visual observations." arXiv preprint arXiv:2206.04779 (2022).

We hope our revisions and our response  address the reviewers’ concerns and further clarify our contributions. If there are any additional questions or comments, we would be happy to address them.

We sincerely thank the reviewer for carefully reading our paper, providing thoughtful responses, and offering valuable feedback. We will certainly incorporate the corresponding discussions and the additional experiments to the final version of the paper.

We remain fully available and would be glad to provide any further details or clarify any remaining points. We believe our responses have addressed most of the concerns raised by you and the other reviewers.

If our replies and the planned revisions have satisfactorily resolved your concerns, we would be grateful if you could kindly reflect this in your evaluation of our work.

Thank you once again for your constructive engagement.