README.md

---
license: apache-2.0
language:
- en
- zh
base_model:
- Qwen/Qwen3-Embedding-4B
pipeline_tag: feature-extraction
library_name: transformers
tags:
- embedding
- retriever
- RAG
- Agentic
- Skills
---
# InsightEmb (Learning Action-Intent Embeddings for Agentic Insight Retrieval)

[![Paper](https://img.shields.io/badge/Paper-arXiv%3A2608.04761-red)](https://arxiv.org/pdf/2608.04761)
[![Code](https://img.shields.io/badge/GitHub-InsightEmb-black?logo=github)](https://github.com/ttchungc/InsightEmb)
[![Project Page](https://img.shields.io/badge/Project-Page-blue)](https://ttchungc.github.io/projects/insightemb/)

**InsightEmb** is an embedding model for **agentic insight retrieval**. It retrieves abstract rules and strategies that resolve an agent's current procedural bottleneck and enable the next useful action.

Unlike conventional dense retrieval, which mainly captures semantic similarity, InsightEmb learns **progress-oriented action-intent matching**. The model is trained on mathematical reasoning data and transfers to interactive agent environments without environment-specific retriever fine-tuning.


For the complete training pipeline, insight-generation scripts, and end-to-end agent evaluation, see the [official GitHub repository](https://github.com/ttchungc/InsightEmb).

---

## ✨ Key Features

- **Progress-Oriented Retrieval**  
  Retrieves insights according to whether they can resolve the current bottleneck and advance the agent toward its goal.

- **Action-Intent Matching**  
  Connects concrete agent states with abstract, procedurally useful rules rather than relying only on topical overlap.

- **Math-Only Retriever Training**  
  Learns transferable retrieval geometry from mathematical problems and reasoning trajectories, without target-environment retriever optimization.

- **Two-Stage Contrastive Curriculum**
  - **Situation-to-Insight Matching:** aligns problems and intermediate reasoning states with abstract heuristic rules.
  - **Situation-to-Experience Matching:** groups different-looking situations that require the same underlying strategy.

- **Partial-Trajectory Supervision**  
  Uses truncated reasoning traces to train retrieval from intermediate states, improving sensitivity to evolving procedural bottlenecks.

- **Cross-Domain Transfer**  
  Evaluated on ALFWorld, WebShop, ScienceWorld, and SRA-Bench, with consistent gains over the base embedder and strong reasoning-oriented retrievers.

---

## 🚀 Quick Start

### Installation

```bash
pip install torch transformers>=4.53.0
```

### 1) Load the Model

Replace `YOUR_HF_NAMESPACE/InsightEmb` with the model repository ID.

```python
import torch
import torch.nn.functional as F
from transformers import AutoModel, AutoTokenizer

MODEL_ID = "YOUR_HF_NAMESPACE/InsightEmb"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# InsightEmb is initialized from Qwen3-Embedding-4B.
tokenizer = AutoTokenizer.from_pretrained(
    "Qwen/Qwen3-Embedding-4B",
    trust_remote_code=True,
    padding_side="left",
)

model = AutoModel.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
).to(DEVICE).eval()
```

### 2) Encode Text

InsightEmb uses last-token pooling followed by L2 normalization.

```python
def last_token_pool(last_hidden_states, attention_mask):
    left_padding = attention_mask[:, -1].sum() == attention_mask.shape[0]
    if left_padding:
        return last_hidden_states[:, -1]

    sequence_lengths = attention_mask.sum(dim=1) - 1
    batch_size = last_hidden_states.shape[0]
    return last_hidden_states[
        torch.arange(batch_size, device=last_hidden_states.device),
        sequence_lengths,
    ]


@torch.inference_mode()
def encode(texts, max_length=8192):
    inputs = tokenizer(
        texts,
        padding=True,
        truncation=True,
        max_length=max_length,
        return_tensors="pt",
    ).to(DEVICE)

    outputs = model(**inputs)
    embeddings = last_token_pool(
        outputs.last_hidden_state,
        inputs["attention_mask"],
    )
    return F.normalize(embeddings, p=2, dim=1)
```

### 3) Retrieve Insights for an Agent State

Queries should include a retrieval instruction. Candidate insights are encoded directly.

```python
retrieval_instruction = (
    "Given the current agent state, retrieve insights that resolve the "
    "current bottleneck and help the agent make progress toward its goal"
)

agent_state = """
Goal: buy a light-grey dining set under $250.
History: searched twice and browsed several result pages.
Observation: many partially matching products are visible, but none selected.
"""

query = f"Instruct: {retrieval_instruction}\nQuery:{agent_state}"

insights = [
    "Keep refining the query until an exact product appears.",
    "Include all critical attributes in the search, then select the required variant before purchasing.",
    "Browse more result pages before committing to a product.",
]

query_embedding = encode([query])
insight_embeddings = encode(insights)

scores = (query_embedding @ insight_embeddings.T).squeeze(0)
top_indices = torch.topk(scores, k=2).indices.tolist()

for rank, index in enumerate(top_indices, start=1):
    print(f"Insight {rank}: {insights[index]}")
```

For dynamic retrieval after every environment step, see the [ALFWorld, WebShop, and ScienceWorld implementations](https://github.com/ttchungc/InsightEmb).

---

## 🎯 Intended Use

InsightEmb is designed for retrieval settings where relevance is primarily **procedural or structural**, including:

- agent memory and insight retrieval
- state-aware retrieval for interactive agents
- skill and strategy retrieval
- reasoning-intensive retrieval
- retrieval of abstract rules with limited lexical overlap

It can be used as a bi-encoder retriever or as the dense component of a hybrid retrieval pipeline.

---

## 📜 Citation



If you find this work useful, please cite:

```bibtex
@misc{chung2026insightemblearningactionintentembeddings,
      title={InsightEmb: Learning Action-Intent Embeddings for Agentic Insight Retrieval}, 
      author={Tsz Ting Chung and Jiangnan Li and Jie Zhou and Mo Yu},
      year={2026},
      eprint={2608.04761},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2608.04761}, 
}
```