Rethinking Latent Representations in Behavior Cloning: An Information Bottleneck Approach for Robot Manipulation

Behavioral Cloning (BC), first introduced by (Pomerleau, 1991), is a well-known Imitation Learning (IL) algorithm that learns a policy by directly minimizing the discrepancy between the agent’s actions and those of the expert in the demonstration data. To learn more generalizable representations, one class of visual representation learning methods pre-train on large video datasets of robotics or humans, enabling rapid application of the pre-trained encoder to downstream robotic tasks. Notable examples include VC-1 (Majumdar et al., 2023), R3M (Nair et al., 2023), and Voltron (Karamcheti et al., 2023) . Meanwhile, another line of research focuses on training on even more extensive and diverse datasets with larger models, such as Internet-scale visual question answering and robot trajectory data (Brohan et al., 2023), as well as a vast collection of Internet videos (Cheang et al., 2024). Additionally, some methods further enhance generalization by incorporating additional sources of information. These include inferring textual descriptions based on the robot’s current state (Zawalski et al., 2024), leveraging visual trajectories (Wen et al., 2024) and generated images (Tian et al., 2025), and integrating 3D visual information (Goyal et al., 2023). However, these methods have not deeply analyzed the redundancy in learned latent representations, and most also lack a solid theoretical foundation. Thus we extend the Information Bottleneck (IB) principle to BC, addressing this fundamental gap.

The Information Bottleneck (IB) principle was first proposed in (Tishby et al., 1999) within the context of information theory. Since then, it has been widely applied in deep learning and various downstream tasks to balance the trade-off between representation accuracy and complexity, including classification (Federici et al., 2019), segmentation (Bardera et al., 2009; Lee et al., 2021), and generative tasks (Jeon et al., 2021). In robotics learning, IB has found notable applications in reinforcement learning, where some works maximize the mutual information between the representation and the dynamics or value function, while restricting the information to encourage the encoder to extract only task-relevant features (Kim et al., 2019; Bai et al., 2021; He et al., 2024). In imitation learning, it has been introduced to solve copycat problems (Wen et al., 2020). Different from prior works, we introduce IB into Behavior Cloning to explore and experimentally validate the redundancy in latent representations in robotics. Additionally, we demonstrate its effectiveness through detailed theoretical analyses.

BC can be formulated as the Markov Decision Process (MDP) framework (Torabi et al., 2018), which is often defined without an explicitly specified reward function, to model sequential action generation problems. The concept of rewards is replaced with supervised learning, and the agent learns by mimicking expert actions. Formally, in robot manipulation, the state at each timestep consists of visual observations $o_{t}$, the robot’s proprioceptive state $s_{t}$, and optionally a language instruction $l$. Let $x_{t}=(o_{t},s_{t},l)$ represent the overall state. The policy $\pi$ maps a sequence of states to an action: $\hat{a}_{t}=\pi(x_{t-\tau:t})$, where $\tau$ indicates the length of the state history. For simplicity, we set $\tau=1$. The optimization process can be formulated as:

where $\mathcal{D}_{e}$ is expert trajectory dataset and $a_{t}$ is action labels. In vanilla BC, $\mathcal{L}$ typically represents the mean squared error (MSE) loss function for continuous action spaces, or cross-entropy (CE) loss for discrete action spaces. In this study, we adopt the continuous action spaces with MSE loss:

Building on this vanilla BC loss, some methods also introduce alignment loss (Jang et al., 2022; Ma et al., 2024) and reconstruction loss (Radosavovic et al., 2023; Karamcheti et al., 2023). However, in this study, to more clearly illustrate the relationship with representation redundancy, we focus solely on the vanilla BC loss.

Estimating mutual information between variables directly is challenging, thus we use Mutual Information Neural Estimation (MINE) (Belghazi et al., 2018) to estimate it. MINE is based on neural networks, which can efficiently handle high-dimensional, continuous, discrete, and hybrid data types without requiring assumptions about the underlying distributions. MINE estimates mutual information by training a classifier to differentiate between samples from the joint distribution $P_{XZ}$ and the product of the marginal distributions $P_{X}\otimes P_{Z}$ of the random variables $X$ and $Z$. MINE uses a lower bound for mutual information based on the Donsker-Varadhan representation (Donsker & Varadhan, 1983) of the Kullback-Leibler (KL) divergence:

where $T_{\theta}:\mathcal{X}\times\mathcal{Z}\to\mathcal{R}$ is a discriminator function modeled by a neural network with parameters $\theta$. We empirically sample from $P_{XZ}$, and for $P_{X}\otimes P_{Z}$, we shuffle the samples from the joint distribution along the batch axis.

Before introducing IB, we first need to define the input $X$ and latent representations $Z$. According to current methods in BC for robot manipulation, as discussed in Section 3.1, the input is typically multimodal, meaning it not only includes RGB images but may also incorporate the robot’s proprioceptive state, language instructions, and more. Previous work has shown that proprioceptive states can lead to overfitting (Wang et al., 2024). Additionally, for the sake of convenience in visualizing $I(X,Z)$, we do not treat the image alone as $X$, as done in previous studies. Instead, we use features extracted from all modalities through respective feature extractors as our input $X$, i.e.,

where $\text{Enc}_{(\cdot)}$ denotes the feature extractor of each modality. Then, regarding how to process the input $X$, or how to fuse information from multiple modalities into latent representations $Z$, we categorize the BC methods in robot manipulation based on feature fusion methods into two types: spatial fusion and temporal fusion.

As illustrated in Figure 2 b), spatial fusion involves extracting spatial features from data at a given time step or concatenating features across multiple time steps along the feature dimensions. This approach does not explicitly differentiate between time steps but instead processes the aggregated features as a whole, emphasizing the modeling of inter-feature relationships. The spatial fusion module can be implemented using Multi-Layer Perceptrons (MLPs), Convolutional Neural Networks (CNNs), Spatial Transformers, or even simple concatenation operations. On the other hand, as illustrated in Figure 2 c), temporal fusion integrates input features by capturing dynamic relationships and dependencies across time steps, enabling the modeling of both long-term and short-term temporal dynamics in sequential data. Temporal fusion modules can be implemented using Recurrent Neural Networks (RNNs), Long Short-Term Memory networks (LSTMs), or Temporal Transformers.

The latent representation $Z$, which integrates both spatial and temporal information, is then passed through a policy head to generate actions. Existing policy heads primarily focus on using MLP, Gaussian Mixture Model (GMM), and diffusion-based policy (DP) heads (Chi et al., 2023; Reuss et al., 2024). For simplicity and clearer empirical demonstration, we use an MLP as the policy head.

The Information Bottleneck (IB) principle is an information-theoretic approach aimed at extracting the most relevant information from an input variable $X$ with respect to an output variable, i.e., action $A$. The central idea is to find a compressed representation $Z$ of $X$ that retains the relevant information needed to predict $A$, while discarding irrelevant parts of $X$ that do not contribute to predicting $A$. The relevant information is quantified as the mutual information $I(X;A)$, and the optimal representation $Z$ is the minimal sufficient statistic of $X$ with respect to $A$. In practice, this can be achieved by minimizing a Lagrangian that balances the trade-off between retaining predictive information and compressing the input, which can be formulated as:

where $\beta$ is the Lagrange multiplier that balances the trade-off between the compression ability and the predictive power. Thus Equation 2 can be modified as:

where $z_{t}=F(x_{t})$ and $F(\cdot)$ denotes the fusion module.

We provide a theoretical analysis of our BC-IB objective in Equation 5. Theorem 4.1 and Theorem 4.2 reveal that the generalization error admits an upper bound governed by the mutual information between the input and the latent representation. By minimizing this mutual information, we effectively tighten the upper bound, thereby improving the model’s generalization performance. Theorem 4.3 elucidates the optimization challenges associated with complex input states. When the input $O$ contains a large amount of state information, as depicted in Figure 2, directly minimizing the mutual information between $O$ and the latent representation $Z$ becomes computationally impractical. To address this, the input $O$ is first compressed into an intermediate embedding $X$ via a fusion network $f$, and the mutual information between $X$ and $Z$ is minimized instead. The theorem establishes that, under certain conditions, this intermediate approach can approximate the ideal optimization results, provided that the embedding $X$ sufficiently preserves the essential information from the original input $O$.

Benchmarks. We mainly evaluate BC with IB across two benchmarks, CortexBench (Majumdar et al., 2023) and LIBERO (Liu et al., 2024). CortexBench is a single-task benchmark. For validation, we selected four imitation learning-related simulators, encompassing a total of 14 tasks: Adroit (2 tasks) (Rajeswaran et al., 2018), Meta-World (5 tasks) (Yu et al., 2020), DMControl (5 tasks) (Tassa et al., 2018), and TriFinger (2 tasks) (Wuthrich et al., 2021). During evaluation, the number of validation trajectories is set to 25, 10, 25, and 25, respectively. LIBERO is a language-conditioned multi-task benchmark. For evaluation, we select four suites: LIBERO-Goal (10 tasks), LIBERO-Object (10 tasks), LIBERO-Spatial (10 tasks), and LIBERO-Long (10 tasks), each focusing on the controlled transfer of knowledge related to task goals, objects, spatial information, and long-horizon tasks, respectively. During evaluation, the number of validation trajectories is set to 20.

Baselines. In CortexBench, we evaluate four representation learning models: R3M (Nair et al., 2023), Voltron (Karamcheti et al., 2023), VC-1 (Majumdar et al., 2023), and MPI (Zeng et al., 2024). In line with the original papers, we use the pre-trained versions of these models to facilitate their application to downstream tasks, keeping the image encoders frozen. Additionally, we introduce two full fine-tuning baselines by replacing these encoders with part of uninitialized ResNet-18 (He et al., 2016) and ViT-S (Dosovitskiy, 2021), denoted as ResNet and ViT, respectively. All methods employ the two fusion techniques described in Section 4.1. For spatial fusion, we use an MLP, and for temporal fusion, we utilize a Temporal Transformer. In LIBERO, we implement four vision-language policy networks. One of them uses a spatial fusion approach, which employs ResNet as the image encoder and an MLP as the fusion module, referred to as BC-MLP. The other three use temporal fusion. Following the original paper, we rename them based on the combination of the image encoder and fusion module: BC-RNN, BC-Transformer, and BC-VILT (Liu et al., 2024). The policy head for all methods is fixed as an MLP. Notably, all baselines with IB are referred to as BC+IB.

Implementation. In CortexBench, for four partial fine-tuning methods, we train for 100 epochs on each task using the Adam optimizer with a learning rate of 1e-3, a batch size of 512, and weight decay of 1e-4, with learning rate decay applied using a cosine annealing schedule. For the two full fine-tuning methods, we train for 50 epochs with a learning rate of 1e-4 and a batch size of 256. All other parameters remain unchanged. In LIBERO, we train for 50 epochs using the AdamW optimizer with a learning rate of 1e-4 and a batch size of 64, decayed using a cosine annealing schedule. For BC+IB methods, the model used in MINE consists of a two-layer MLP, with a learning rate of 1e-5. The Lagrange multiplier in Equation 6 ranges from 1e-4 to 5e-3 in this work.

Model Selection. For single-task benchmark CortexBench, we test the model every 5 or 10 epochs and select the model with the highest success rate. For multi-task benchmark LIBERO, we select the model from the final epoch.

The appendix will provide detailed descriptions of each benchmark (Section B.1), all baselines (Section B.2), implementation details (Section B.3), and the rationale behind the model selection (Section B.4).

The Selection of Fusion Method. We first evaluate the effectiveness of the two fusion methods in the baselines on CortexBench, with the results shown in Figure 3.

Finding 1: For simple single-task scenarios, spatial fusion is more efficient and effective than temporal fusion. As shown in Figure 3 b), the performance of methods with temporal fusion drops significantly. From Figure 3 a), this can be attributed to the slower loss reduction in methods using temporal fusion, which results in higher loss at the same training epoch. Therefore, we focus exclusively on presenting the results for methods employing spatial fusion.

Results. We next report the performance, i.e., success rate, of the baselines and baselines with IB on the single-task benchmark CortexBench in Table 1 with a full-shot setting. Based on results, we derive the following findings.

Finding 2: Whether using full fine-tuning or partial fine-tuning, all vanilla BC methods with different visual backbones incorporating IB outperform their vanilla counterparts across the board. In some benchmarks, the improvements are substantial. For example, ResNet with IB achieves a 10.01% improvement on DMControl, and VC-1 with IB shows a 4.80% improvement on Meta-World. In Section C.1, we report the success rate for each task, where significant improvements can be observed in certain tasks.

Finding 3: Finding 2 implicitly suggests that the latent representation $Z$ derived from input $X$ is redundant. Therefore, compressing information from input is essential, which can further enhance performance.

Finding 4: For simple single-task downstream tasks, full fine-tuning of a simple, uninitialized model (ResNet) is sufficient and may even outperform a pre-trained larger model. However, the latter is more efficient for faster fine-tuning and deployment, and proves to be more effective for more complex tasks (Burns et al., 2023).

Results. We report the performance on the multi-task benchmark LIBERO with a full-shot setting in Table 2.

Finding 5: For more complex language-conditioned multi-task scenarios, all baselines with different backbones incorporating IB consistently show performance improvements across all LIBERO benchmarks. For example, BC-VILT achieves large gains of 7.66% and 9.00% on LIBERO-Goal and LIBERO-Object, respectively, while BC-RNN shows a significant improvement of 10.83% on LIBERO-Goal. IB proves to be more effective in more complex environments and settings. We attribute this to the difference in task complexity: in CortexBench, the history length is 3, while in LIBERO, it is 10, with LIBERO being a multi-task benchmark and CortexBench being single-task benchmark. The increased data complexity (task quantity and input information) suggests a higher level of data redundancy, making IB even more effective.

Finding 6: We observe that in complex multi-task scenarios with more intricate inputs, such as a greater number of input modalities and extended historical information, using the Temporal Transformer in temporal fusion proves to be more effective than both spatial fusion and RNN-based temporal fusion. The evidence lies in the fact that the average success rates of BC-Transfomer and BC-VILT are over 30% higher than those of BC-MLP and BC-RNN. This is likely because Temporal Transformers excel in handling long-range interactions and capture dynamic dependencies across time steps, where RNNs and spatial fusion methods may struggle. This finding, together with Finding 1, underscores the specific scenarios in which each fusion method is most applicable.

Finding 7: IB is particularly effective for tasks requiring diverse feature extraction, such as distinguishing distinct task objectives or differentiating between various objects, as in LIBERO-Goal and LIBERO-Object. By filtering out irrelevant information, IB facilitates better generalization and more compact representations. However, its impact is less pronounced in spatial and long-horizon tasks, such as LIBERO-Spatial and LIBERO-Long, which heavily depend on structural and sequential dependencies that may be disrupted by excessive compression.

Effect of the Lagrange multiplier $\beta$ of Equation 6. This experiment evaluates how incorporating IB enhances performance. Since the MINE model’s parameters are fixed, the key difference between BC+IB and BC lies in the parameter $\beta$, which balances compression and predictive power. For LIBERO experiments, $\beta$ is explored within 1e-4, 1e-3, 5e-3, 1e-2. Finding 8: As shown in Figure 5, IB improves performance within a specific $\beta$ range, with a peak observed at an undetermined value. However, across all experiments, $\beta$ around 1e-4 consistently yields stable improvements. The selected $\beta$ values for each experiment are detailed in C.1.

Effect of the Number of Demonstrations. We evaluate IB’s effectiveness in few-shot settings, as few-shot learning is crucial for fine-tuning on domain-specific tasks in real-world applications. Finding 9: As shown in Figure 6, IB significantly improves performance even with limited data, highlighting its effectiveness in real-world scenarios where data is scarce. This further underscores the potential of IB in improving model generalization in practical settings.

Visualizations of $I(X,Z)$. As shown in Figure 4, BC+IB achieves a greater reduction in $I(X,Z)$ compared to vanilla BC, leading to significant performance improvements and further validating the effectiveness of IB. For instance, in LIBERO-Goal, incorporating IB reduces $I(X,Z)$ to one-quarter of its original value, resulting in a 7.7% increase in success rate.

We provide real-world experiments in Section C.2.