Paper: $\pi$2vec: Policy Representation with Successor Features (ICLR-2024)

Reviewer comment (Evaluation):
Why don't you use ranking metrics (Mean Average Precision, DCG, etc)? These seems quite relevant for evaluating methods for offline policy selection. As you mentioned, absolute error in terms of value prediction may not be always relevant.

Author response:
**Ranking metrics in the experiments:** We add additional details about metrics in our general comment and update the Appendix accordingly. To answer the reviewer’s specific question, Correlation is a ranking metric that evaluates the ranking of the estimated policies’ values, which we clarify in the manuscript. It is commonly adopted for Offline Policy Selection [1] and can be used for comparing policies. We report correlation for all our experiments to assess $\pi$2vec’s performances on ranking policies.

Reviewer comment (Experiments):
You decided to focus on visual representation without clearly explaining why. I think it is important to evaluate the approach also on state-based observation. How would you select base state features in this setting?

Author response:
**Visual representations**: Commonly available foundation models are usually trained on visual (e.g., VIT, TAP) and text-vision data (e.g., CLIP). Those features are commonly adopted in robotics tasks [2]. Motivated by the fact that learning RL policies from vision is hard for such a task at hand, we explore how $\pi$2vec for vision-based RL benefits from visual foundation models thanks to their geometrical (TAP), semantical (CLIP), and classification-related (VIT) features. Extending $\pi$2vec for state-based representations requires foundation models that include proprioception or other state-based information as part of their input. Although state-based features are relevant for future work, in this paper, we focused on showing that visual representations are sufficient for policy representations in the domains that we considered.

Reviewer comment (Experiments):
Similarly, I would have expected an evaluation on standard offline benchmarks (e.g., D4RL or ExoRL). For example, if I'm not mistaken, ExORL provides also image-based observations.

Author response:
**Evaluation on standard offline benchmarks**: We are interested in vision-based robotic environments, and we leave further evaluation of $\pi$2vec on different scenarios as future work. [3] is limited in scope to synthetic environments with a focus on state-based RL (e.g., CartPole, Walker), while [4] evaluates a broader range of environments, including synthetic (Maze2D), navigation (CARLA, AntMaze), and robotic manipulation (Kitchen) environments. In our experiments, we included the Kitchen environment along with other real and synthetic environments for vision-based robotics (Metaworld, RGBStacking, Insert gear sim and real)

Reviewer comment (Theory):
Limited theoretical analysis: The paper shows empirically that the obtain successor representations are effective, but provides limited insight/intuition or analysis into why combining successor features and foundation models results in good policy embeddings.

Author response:
**Limit analysis on why combining successor features and foundation models is effective:**
$\pi$2vec is theoretically justified by the successor feature framework [1]. We based the method on the assumption that visual foundation models consistently represent pertinent features of the environment and are adaptable as reward predictors[2]. Our focus is, therefore, on the empirical evaluation of $\pi$2vec’s representations. 
We validate this assumption in paragraph (iii) of the experimental Section when we validate $\pi$2vec for fully offline policy evaluation. In this setting, we fit the reward predictor $\hat{r}=\langle\phi(s), w_\text{task}\rangle$ from offline data, where $\phi$ is a foundation model, and we estimate $V^\pi(s)$ by injecting $w_\text{task}$ in Equation 1 of the main manuscript. Even if this assumption might not be true in general, we show in Table 3 that it is reasonable for the settings we consider. Table 3 reports that estimating value functions with $\pi$2vec representations is a strong offline policy evaluation method compared to FQE, a state-of-the-art method for OPE.

Reviewer comment (Experiments):
Restricted to offline setting: The method seems to require pre-collected offline datasets and cannot be applied in an online setting where policies interact with the environment. The offline assumption limits the applicability.

Author response:
Dear reviewer *eE8e*, thank you for the comments and suggestions to improve our paper. We address the question about online representation in the general comment.

Reviewer comment (Evaluation):
Why do successor features plus foundation models work well for policy representation? More analysis on why this combination is effective compared to other representations would be useful.

Author response:
**Comparing feature representations for $\pi$2vec:**
We clarify our insights about $\pi$2vec representations. We show that $\pi$2vec with foundation models always outperforms random features (see Tab. 1 and 2). We compare geometrical (TAP) and semantical (CLIP) features in Table 1 for the task Insert Gear (sim). We show that geometrical representations improve for correlation, while semantical task helps for Regret@1. Furthermore, we discuss our insight by comparing CLIP and VIT for policy representations in paragraph (iv) of the Appendix. From our analysis, we suggest that $\pi$2vec with CLIP features performs better when the historical policies do not solve tasks effectively. This is often the case when starting working on a new robotic task, e.g., insert gear (real). On the other hand, VIT features help find the best policy among a set of well-performing policies, which is ideal in ablating policies that achieve good performances for the task.

Reviewer comment (Experiments):
Can this method be extended to an online setting where policies interact with the environment?

Author response:
Dear reviewer *eE8e*, thank you for the comments and suggestions to improve our paper. We address the question about online representation in the general comment.

Reviewer comment (Evaluation):
How does dataset composition impact the quality of the policy embeddings? Is there any further analysis on dataset requirements and relationships between dataset and representation quality?

Author response:
**Dataset quality**: we explore dataset quality in paragraph (v) in Section 7.4 of the Appendix. We collect an enhanced dataset of trajectories for all Metaworld tasks and point-of-views (12 settings). A dataset of trajectories better resembles the behavior of historical policies that are used for $\pi$2vec’s representations. We report results in Table 5, showing that $\pi$2vec improves performance when adopting an improved dataset. Similar findings hold when adopting $\pi$2vec on online data.

Reviewer comment (Experiments):
Have you considered any other methods to aggregate the state-dependent successor features? Averaging seems effective but overly simplistic; are there alternatives that may capture policies better?

Author response:
Other aggregation methods: we thank the reviewer for the suggestion. In this initial work, we focus on proposing $\pi$2vec’s representations, and we leave aggregation strategies as further work.

Reviewer comment (Experiments):
Potentially missing baseline: I am not closely follow the OPE literature, but could the authors explain why [1] is not suitable for this setting? We should be able to get ranking results by that approach.

Author response:
**Comment on [1]**: Irpan et al. [1] propose an OPE approach for robotics. Authors assume that reward is sparse and binary (success/failure in an episode), and they cast OPE as a binary classification problem. Moreover, their assumption about sparse and binary rewards does not hold in our setting, as the environments that we adopted have continuous and dense rewards. In experiment (iii) of the main paper, we compare our method to FQE, which is a strong baseline in  OPE (similar to [1] but without the assumption of binary rewards), showing promising results.

Reviewer comment (Experiments):
Missing key ablations: I think there are implicit key assumptions such as 1) the historical policies' performance should most likely cover the unseen policy. If not, and the unseen policy's performance is way better, I hardly expect this approach would work.

Author response:
Coverage of historical policies: While it is true that very differently performing policies are unlikely to provide useful information for policy evaluation of the new policies, we have some experiments that look at the question of how the difference in the policies affects the performance of the method. We explore the point raised by the reviewer in paragraph (vi) in the Appendix. We adopt $\pi$2vec for representing intermediate checkpoints for a set of policies for Metaworld assembly task, and we test on held-out policies. Intermediate checkpoints for the same policy are similar to each other and different from checkpoints of other policies, especially at the beginning of training [2].  $\pi$2vec greatly outperforms actions in this setting for all metrics, proving its robustness to policy coverage.

Reviewer comment (Experiments):
Missing key ablations: I think there are implicit key assumptions such as 2) the canonical states coverage. How does it affect the performance if state coverage of MDP is small? More specifically, if the historical dataset only covers 50% of states that the unseen policy would visit in the MDP, how does it affect the performance? It would be good to have a comprehensive understanding when this approach can be effective.

Author response:
**Canonical state coverage**: We sample canonical states uniformly from the dataset that was used for training offline RL policies. Even though selecting canonical states to represent the environment better is intuitively beneficial, even simple sampling at random worked well in our experiments. We conduct further experiments to ablate the state coverage. We adopt demonstrations from Metaworld Assembly task and adopt the initial state of each trajectory for $\pi$2vec’s and Action’s representation. By using the initial state of a trajectory, $\pi$2vec cannot rely on the state coverage. We report the results in the following table. We show that $\pi$2vec is robust to state coverage, showing SOTA performances even when only initial states are used. We add this experiment in the Appendix of the revised manuscript.

*Assembly Left*
| Feature  	|	NMAE | Correlation | Regret@1 |
|--------------|---------|-------------|----------   |
| CLIP     	|   0.381 |   	0.359 |	0.287 |
| $\Delta$CLIP | **0.260** | 	**0.592** |  **0.087** |
| Random   	|   0.422 |   	0.252 | 	0.46 |
| VIT      	|   0.366 |    	0.26 |	0.347 |
| Actions  	|   0.356 |   	0.503 |	0.222 |

*Assembly Right*
| Feature   |	NMAE | Correlation | Regret@1 |
|-----------|---------|-------------|----------   |
| CLIP  	|   0.363 |   	0.023 |	0.365 |
| $\Delta$CLIP      	| **0.242** | 	**0.582** |  **0.096** |
| Random	|   0.334 |   	0.313 |	0.212 |
| VIT   	|	0.27 |   	0.345 |	0.304 |
| Actions   |   0.405 |   	0.369 |	0.263 |

*Assembly Top*
| Feature  	|	NMAE | Correlation | Regret@1 |
|--------------|---------|-------------|----------   |
| CLIP     	|   0.463 |   	0.270 |	0.330 |
| $\Delta$CLIP | **0.305** | 	**0.594** |  **0.078** |
| Random   	|   0.394 |   	0.277 |	0.328 |
| VIT      	|   0.418 |   	0.020 |	0.417 |
| Actions  	|   0.414 |   	0.554 |	0.106 |

Reviewer comment (Theory):
In section 3.4, \Phi_\pi represents an aggregated average effect of the behavior of π, I think this highly relies on the implicit assumptions mentioned above. It would be good that authors can clarify.

Author response:
**Interpretation for $\Phi_\pi$**: As mentioned by reviewer *4YPE*, $\Phi_\pi$ represents policy $\pi$ by aggregating the successor feature representation for policy $\pi$ over a canonical set of states. $\pi$2vec performances are influenced by (i) historical policy coverage, (ii) coverage of the training data, and (iii) canonical states. Our experiments show that (i) $\pi$2vec is robust to historical policy coverage (paragraph *vi* and *x*, Section 7.4 in the revised Appendix), (ii) including online trajectories in the training set improves $\pi$2vec’s representations (paragraph *v*, section 7.4 in Appendix), and (iii) canonical state coverage has a limited impact on $\pi$2vec (paragraph ix in the revised manuscript in the Appendix). We hope our answers address the question and are open for further discussion.

Reviewer comment (Evaluation):
This paper needs more analysis regarding why certain features work better for certain settings. At the moment, if I implement this idea, I would have to enumerate all possible pairs.

Author response:
**Better analysis of the features**: we present our intuition regarding the features to adopt in paragraph (iv) in the Appendix. We compare CLIP and VIT across all environments. Tables 1 and 2 show a correlation between the choice of the feature encoder and policies’ true return values. CLIP performs better than VIT in Insert Gear (Sim), Insert Gear (Real), and Metaworld. Table 4 reports policy performances in these environments, showing they are the hardest to solve, as the standard deviation among policies is high. These results demonstrate that CLIP is robust when the task is hard, and most of the policies fail to solve it. We stress that $\pi$2vec might be especially useful for this use case, as corresponds to a scenario where there is an ongoing research effort to train policies for a particular task on a real robotic platform. On the other hand, VIT is the best option when the standard deviation of policy returns are low or negligible (e.g., Kitchen and RGB stacking), making this feature most suited for ranking policies that consistently solve the task.

Reviewer comment (Experiments):
The authors only compare to 1 baseline algorithm that only depends on the actions of the policies in the offline dataset. Meanwhile, the proposed method uses the state-action in the offline dataset to learn the successor feature components. I don't know if the performance increase of the baseline is due to the proposed idea or the additional information the proposed idea has access to.

Author response:
**Baseline comparison**: We point out that the action representations are extracted by running the historical policies on the set of canonical states to compute their actions. We adopt the same set of canonical states for $\pi$2vec and action representations. In this regard, the baseline and $\pi$2vec have access to the same data, and the baseline uses state information implicitly by deciding which actions should be compared. Active Offline Policy Selection (AOPS) [1] stands alone as a notable work that delves into policy representation from offline data with the task of deciding which policies should be evaluated in priority to gain the most information about the system. We argue that this is a strong baseline, as demonstrated in [1] "actions baseline" proves to be insightful for understanding policy behavior (Figure 3 in [1]) and predicting policy performance (Figure 5 in [1]).  We are not aware of any previous work that addresses this particular problem by bringing foundation model representations into offline policy representation and, in particular, any other prior method that would use states in a more explicit way.

Reviewer comment (Experiments):
Why is action representation the only baseline applicable baseline? The underlying data has states and actions. Would it be unfair to condition a baseline on both state-action pairs?

Author response:
**Baseline comparison**: We point out that the action representations are extracted by running the historical policies on the set of canonical states to compute their actions. We adopt the same set of canonical states for $\pi$2vec and action representations. In this regard, the baseline and $\pi$2vec have access to the same data, and the baseline uses state information implicitly by deciding which actions should be compared. Active Offline Policy Selection (AOPS) [1] stands alone as a notable work that delves into policy representation from offline data with the task of deciding which policies should be evaluated in priority to gain the most information about the system. We argue that this is a strong baseline, as demonstrated in [1] "actions baseline" proves to be insightful for understanding policy behavior (Figure 3 in [1]) and predicting policy performance (Figure 5 in [1]).  We are not aware of any previous work that addresses this particular problem by bringing foundation model representations into offline policy representation and, in particular, any other prior method that would use states in a more explicit way.