We propose ToggleMimic, an end-to-end imitation learning framework that generates robotic motions from textual instructions, enabling language-driven multi-task control. The algorithmic framework is illustrated in
Figure 1. The teacher policy (Fused Attention Policy, FAP) is trained using Proximal Policy Optimization (PPO) [
46], utilizing proprioceptive, privileged data (including friction, dataset trajectory, etc.), and text labels as observation data. The teacher policy trains an action generation model incorporating privileged data, which serves as the expert strategy for high-quality action generation in the second stage.
The distillation process explicitly operates without privileged data, producing a policy requiring only text labels and robot state. This Student Distilled Policy (SDP) enables effective reasoning in real-world deployment through its minimal-input design.
3.2. Fused Attention Policy (FAP)
Our Fused Attention Policy (FAP) trains a strategy that tracks motion using motion trajectories from AMASS dataset, while enabling action switching via text commands. The FAP policy’s input is divided into two components: (1) a combined observation of proprioceptive and privileged data; (2) 512-dimensional text embeddings derived from the pre-trained XCLIP model.
These two components jointly enable the robot to align its motion trajectory with text commands. The input structure for FAP’s observations is detailed in
Table 1.
Our FAP network features three core designs: the Early Fusion (EF) layer, the Observation feature extraction module ( module), and the LightCrossAttention module (LCA module).
- (1)
Early Fusion (EF) layer
The Early Fusion (EF) layer fuses the normalized observation with the text embedding via an MLP, producing a text-augmented state representation that modulates the proprioceptive using text semantics.
where
—text embedding;
—proprioceptive observation;
—privileged observation.
The EF layer provides contextual constraints for the distribution of state features under different task instructions, enabling the network to dynamically adjust the state encoding strategy based on textual semantics. For instance, for instructions with significant semantic differences such as “raising a hand”, the fused features will vary along different subspace directions. This provides differentiated semantic guidance for subsequent motion feature extraction and cross-modal attention mechanisms. This low-level conditional feature fusion mechanism introduces semantically relevant constraints at the state feature level, facilitating the generation of more stable and semantically consistent motion trajectories by the model during language-driven action transitions.
- (2)
Observation feature extraction module ( module)
The Observation feature extraction module (
Figure 2) is a core component of the FAP network, responsible for extracting state-semantic representation from the text-augmented state representation that can be used for action generation.
where
—LSTM, and
—a two-layer ReLU-MLP.
To capture temporal continuity in robotic motion while extracting nonlinear features from text-augmented state representation (generated in the EA layer), the module employs a hybrid LSTM-MLP structure. This design enables simultaneous nonlinear feature extraction from current observations (via MLP) and temporal memory for historical state sequences (via LSTM), ensuring temporal coherence across the motion trajectory. Consequently, the policy network achieves balanced local responsiveness and global dynamics during motion generation, producing stable state-semantic representations for subsequent multimodal action generation.
- (3)
LightCrossAttention (LCA)
To enhance model interpretability for text-guided robotic actions, we introduce LightCrossAttention (LCA) (
Figure 3), a lightweight cross-modal attention module. LCA adopts a cross-attention structure that learns semantic-action correspondence, allowing the policy network to generate actions directly conditioned on linguistic instructions. By incorporating residual fusion and linear mapping, LCA maintains the stability of representation learning while minimizing computational overhead, enabling efficient deployment on resource-constrained robotic systems.
Given the state-semantic representation and text embeddings, the LCA module aligns the text features to the state representation space through a learned mapping, enabling semantic correspondence between language and sensory input.
where
is the ReLU activation function,
and
are linear projection matrices,
represents the dimension of the hidden space, and
Y represents the output semantic-aligned features.
The multi-head attention computes the association weights between the state-semantic representation and the text embedding, and the output attention result A represents the weighted semantic response of the state-semantic representation feature at each time step in the text embedding vector space. To enhance stability, residual connections and layer normalization are used for fusion. Finally, a linear mapping transforms the fused cross-modal features into semantic-aligned features.
- (4)
Action Generation Layer
After completing the cross-modal feature interaction, the model fuses the state-semantic representation features with the semantic-aligned features.
where
—state-semantic representation and
Y—semantic-aligned features.
This element-wise additive fusion retains the continuous dynamic information of the two types of state features, enabling the model to simultaneously consider the semantic constraints of the robot’s current state and the textual instruction during the action generation phase. The fused features,
, are then fed into the Action Head, which maps the high-dimensional semantic features to the robot’s action space via an MLP:
where
—robot action space.
This structure achieves the fusion of state-semantic representation and semantic-aligned features, directly mapping them to the action space. It enables the model to generate control commands that comply with both task semantics and kinematic constraints, under the guidance of a joint semantic representation.
To accelerate the teacher policy’s learning from the dataset, our FAP strategy is trained using PPO algorithms. The reward function and its weights are shown in
Table 2.
3.3. Student Distilled Policy (SDP)
Our Student Distilled Policy (SDP) replaces privileged inputs in the Fused Attention Policy, eliminating the need for inaccessible observations during deployment. To accelerate distillation efficiency, SDP preserves the FAP’s core network architecture. By inheriting this structure, the student avoids redundant learning of sensor-to-state mappings and instead focuses on reconstructing the decision-making features and control distributions of the Teacher within a constrained network. This significantly improves distillation stability and convergence speed. Crucially, we incorporate three domain-specific enhancements for robust generalization:
To maintain semantically consistent control features after input data compression, the SDP introduces a Gating Network module after the Action Generation Layer. The state-semantic representation
, obtained through feature extraction, and the semantic-alignment feature
Y are combined by computing fusion weights using a gating unit.
where
—Sigmoid function and ⊙—element-wise multiplication.
The core mechanism employs learnable gating parameters to dynamically modulate the fusion of state-semantic representation and semantic-aligned features in the SDP. This enables the network to autonomously adjust the dependency strength between textual representations and proprioceptive across varying conditions. The gating structure allows the SDP to adaptively learn feature dependency dynamics during training—with robustness to observation noise or representation redundancy—ensuring action stability. Furthermore, the gating mechanism enhances SDP’s capacity for novel textual input recognition, strengthening task-switching capability and generalization across diverse operational scenarios.
- (2)
Action Classification Head (ACH)
To enhance distributional stability and action-space consistency in the student policy during action generation, we introduce an Action Classification Head (ACH) after the action distribution output of the Action Generation Layer. This module enables the model to distinguish between action categories aligned with instructional text while generating action distributions, thereby achieving separated boundary representations in the latent action space.
where
: the classification weight matrix,
: the action output distribution (SDP),
: the number of action categories.
Meanwhile, the action classification auxiliary head employs Cross-Entropy Loss:
where
: action class label (one-hot encoded) at time
t, and
: predicted probability.
Our Action Classification Head explicitly constrains the policy network to maintain distributional stability within the shared latent action space, making the distributions of similar actions more concentrated. It provides action-space-aligned gradient supervision during optimization, thereby enhancing generalization across instructional text and distributional stability in the initial training phase.
- (3)
Parameterized Distribution
In SDP, action generation is modeled as a Gaussian distribution in the latent action space, with the action distribution
output by the backbone network and the standard deviation
as a learnable positive parameter.
where
: action distribution, which is output by the backbone network;
: standard deviation, used to characterize the uncertainty of the policy.
Our method employs a learnable global parameter vector, assigning a shared standard deviation to each action dimension:
where ⊗: the broadcast replication operation;
: action space
The Gaussian distribution in the latent action space provides a stable continuous gradient path for loss function computation. This formulation models action uncertainty, accelerating convergence during the initial training phase while maintaining a smooth policy distribution within the latent policy space in subsequent training stages.
- (4)
Loss function
Apart from the Action Classification Head, the loss function of SDP consists of two additional parts:
Behavioral Cloning loss
:
To align the Student Distilled Policy with the Fused Attention Policy during distillation while minimizing cumulative errors, we introduce a behavioral cloning constraint based on Dataset Aggregation (DAgger) [
49]. The behavioral cloning loss minimizes the discrepancy between the mean action output by the SDP and the FAP, allowing the student to gradually approximate the teacher’s behavior patterns in the policy space. This ensures consistency in basic movements. Meanwhile, DAgger’s data aggregation mechanism further reduces cumulative errors caused by distribution shifts through dynamic sampling of teacher trajectories and student behavior trajectories, thereby ensuring the stability and consistency of the distillation process.
Bidirectional KL divergence loss
:
To enhance the distributional stability and generalization of the student policy within the latent action space, we introduce a bidirectional KL divergence loss during distillation. The forward KL divergence, , encourages the student to align its action distribution with high-probability regions of the teacher’s distribution in the latent action space, ensuring coverage of diverse behavioral modes. Conversely, the reverse KL divergence, , suppresses arbitrary dispersal of the student distribution into low-probability regions of the teacher, preventing distribution collapse and stabilizing exploration. This symmetric constraint mitigates instability from distributional variance mismatch while maintaining action diversity, enabling efficient distribution convergence during distillation.