# TaskMem
**Repository Path**: ByteDance-Seed/TaskMem
## Basic Information
- **Project Name**: TaskMem
- **Description**: No description available
- **Primary Language**: Unknown
- **License**: Apache-2.0
- **Default Branch**: main
- **Homepage**: None
- **GVP Project**: No
## Statistics
- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-09-15
- **Last Updated**: 2026-09-17
## Categories & Tags
**Categories**: Uncategorized
**Tags**: None
## README
# Task-Focused Memorization for Multimodal Agents
[arXiv](https://arxiv.org/pdf/2605.31075) | [Demo](https://hope-rita.github.io/TaskMem-demo/) | [Model](https://huggingface.co/ByteDance-Seed/TaskMem) | [Phase One](https://github.com/Hope-Rita/TaskMem-PhaseOne)
> This repository is the **inference framework** of TaskMem: it streams a
> video clip-by-clip, builds episodic long-term memory, and runs
> streaming-memory QA evaluation against an off-the-shelf VLM. Training of
> the base memorization policy is released as the companion repository
> [**TaskMem-PhaseOne**](https://github.com/Hope-Rita/TaskMem-PhaseOne).

> **(a) TaskMem Architecture.** An agent receives streaming multimodal
> inputs and streaming tasks from the environment, and continually updates
> its long-term memory through a learned **memorization policy**.
> **(b) Running Example.** Tasks arriving at different time steps shift
> the focus of the memorization policy
> $\pi_{t_i} \to \pi_{t_j} \to \pi_{t_k}$, so the same scene is
> memorized in different ways depending on what the agent will be asked.
## Abstract
Long-term memory is essential for multimodal agents to build coherent experience, accumulate world knowledge, and achieve continual learning. However, constructing effective memory goes beyond memory module design and basic requirements such as accuracy and fidelity; the key challenge lies in determining **what to memorize**. Multimodal agents, such as embodied agents, continuously perceive, reason, and act in real or virtual environments, receiving an unbounded stream of multimodal observations. From this combinatorial explosion of information, an agent must selectively retain content that is relevant to its role in the environment and valuable for future tasks. To bridge this gap, we frame memory generation as a learnable memorization policy and introduce **TaskMem (Task-focused Memorization Policy Learning)**, a reinforcement-learning-based framework that enables the policy to dynamically adjust its focus to the demands of real tasks encountered in the environment. TaskMem adopts a two-phase training paradigm: **Phase One** learns *how* to memorize by optimizing memory quality under fundamental fidelity requirements; **Phase Two** occurs after deployment, where the agent learns *what* to memorize by tuning an adapter on its base MLLM, using recent environment tasks to define a reward model that guides the memorization policy toward task-relevant content. To evaluate our approach, we reformulate VideoMME, EgoLife, and EgoTempo into **streaming benchmarks** that simulate a realistic setting in which an agent processes streaming observations and handles tasks arriving online. To isolate memory assessment, the questions must be answered using only the agent's memory, without access to raw video. Built on **Qwen3-VL-30B-A3B**, TaskMem improves VQA accuracy by **6.3%**, **7.0%**, and **5.3%** on these benchmarks, respectively.
## Method

> **Phase One** trains the memorization policy on a frozen video dataset
> against fidelity-style rewards — **Format**, **Thinking Length**,
> **Quality** and **Richness** — to learn *how* to memorize.
> A subsequent task-relevance phase adapts the policy to specific
> deployment tasks; see the [paper](https://arxiv.org/pdf/2605.31075) for
> details.
## Experimental Results
| Method |
Video-MME |
EgoLife |
EgoTempo |
| Acc. | Cov. | Prec. |
Acc. | Cov. | Prec. |
Acc. | Cov. | Prec. |
| EgoGPT | 44.3 | 58.7 | 75.5 | 19.2 | 28.2 | 68.1 | 15.0 | 33.5 | 44.9 |
| HippoMM | 48.9 | 66.6 | 73.5 | 30.4 | 43.4 | 70.0 | 15.8 | 30.8 | 51.1 |
| M3-Agent | 62.5 | 77.7 | 80.4 | 21.8 | 30.8 | 70.8 | 16.0 | 36.3 | 44.2 |
| Gemini-1.5-Pro | 55.3 | 65.9 | 83.9 | 39.4 | 51.6 | 76.4 | 19.7 | 34.3 | 57.4 |
| Gemini-2.5-Pro | 63.2 | 74.8 | 84.4 | 43.8 | 56.6 | 77.4 | 25.8 | 42.3 | 61.0 |
| GPT-5.2 | 67.3 | 80.8 | 83.3 | 34.8 | 48.2 | 72.2 | 32.1 | 51.4 | 62.4 |
| Qwen3-VL-30B-A3B | 61.6 | 74.7 | 82.5 | 38.4 | 52.4 | 73.3 | 22.3 | 38.9 | 57.2 |
| TaskMem (Ours) | 67.9 | 79.3 | 85.6 | 45.4 | 56.4 | 80.5 | 27.6 | 43.7 | 63.2 |
See the [project page](https://hope-rita.github.io/TaskMem-demo/) for the
full table and a live demo.
## Run Locally (Inference)
This repository runs the **memorization pipeline at inference time**: it
takes a raw video, streams it clip-by-clip, and builds an episodic
long-term memory that downstream QA can consume. The episodic step is
driven by an off-the-shelf VLM (Gemini / GPT API or vanilla Qwen3-VL via
local vLLM), or by our released **Phase-One** TaskMem checkpoint on
Hugging Face ([`ByteDance-Seed/TaskMem`](https://huggingface.co/ByteDance-Seed/TaskMem)),
loadable through the same `qwen3_vl_vllm` backend by passing it as
`--episodic_model_path`.
### 1. Install
```bash
git clone https://github.com/ByteDance-Seed/TaskMem.git
cd TaskMem
bash setup.sh
```
`bash setup.sh` installs the Python deps (`moviepy`, `pydub`,
`hdbscan`, `insightface`, `json-repair`, `openai`) plus the `ffmpeg`
system package. The Gemini / GPT default path does not need vLLM.
If you want to use the optional local Qwen3-VL episodic backend
(`qwen3_vl_vllm`), install `vllm` and `qwen-vl-utils` separately.
### 2. Configure the LLM judge
Copy the template, fill in your credentials, and point
`TASKMEM_API_CONFIG` at it:
```bash
cp configs/api_config.json configs/api_config.local.json
export TASKMEM_API_CONFIG=configs/api_config.local.json
```
Required because the default ASR / voice / episodic backends are all
Gemini, and because all QA evaluation scripts call an LLM judge.
### 3. Run the pipeline
End-to-end on a single video:
```bash
export VIDEO_PATH=./data/videos/demo.mp4
export OUTPUT_FOLDER=./out/demo
# Default: Gemini for ASR / voice / episodic. API-only, no GPU required;
# expects TASKMEM_API_CONFIG to be set.
bash examples/run_baseline.sh
# Vanilla Qwen3-VL (local vLLM) for the episodic step.
EPISODIC_MODEL=qwen3_vl_vllm \
EPISODIC_MODEL_PATH=Qwen/Qwen3-VL-30B-A3B-Thinking \
bash examples/run_baseline.sh
# Released Phase-One TaskMem checkpoint (same vLLM backend).
EPISODIC_MODEL=qwen3_vl_vllm \
EPISODIC_MODEL_PATH=ByteDance-Seed/TaskMem \
bash examples/run_baseline.sh
```
The script runs three `src/main.py` invocations: (1) `--process_audio`
for ASR + diarization, (2) `--process_video --process_voice` for face
detection, speaker matching, and per-clip mp4 rendering, and (3)
`--generate_episodic` to turn the per-clip context into episodic memory.
After it finishes you get `$OUTPUT_FOLDER/_.pkl`,
a `LongTermMemory` object whose `to_string("episodic")` renders the
episodic memory text consumed by the QA scripts below.
### 4. Question Answering and Evaluation
Every QA script loads the judge LLM credentials via
`TASKMEM_API_CONFIG`. The `--memory_tag` / `--memory_name` flag must
match the `WRITE_TAG` used during memory generation (`baseline_ep` by
default).
```bash
# Video-MME (multiple-choice)
python test/test_videomme_qa.py \
--memory_folder ./out/videomme \
--video_info_root ./data/videomme/info \
--memory_name baseline_ep --memory_type episodic
# EgoLife (multiple-choice, multi-day egocentric)
python test/test_egolife_qa.py \
--qa_file ./data/egolife_qa.json \
--memory_root ./out/egolife \
--memory_tag baseline_ep --memory_type episodic \
--output_file ./out/results/egolife_baseline.jsonl
# EgoTempo (free-form, LLM-judged)
python test/test_egotempo_qa.py \
--memory_folder ./out/egotempo \
--video_info ./data/egotempo/video_info.json \
--memory_name baseline_ep --memory_type episodic
```
### Backends
- **ASR / voice / episodic models** accept any `gemini-*` or `gpt-*`
key from your `TASKMEM_API_CONFIG`, or `qwen3_vl_vllm` for local
Qwen3-VL via vLLM (install `vllm` and `qwen-vl-utils` separately).
Run `python src/main.py --help` for the full flag list.
- **Face / voiceprint backends.** Face detection + clustering use
[insightface](https://github.com/deepinsight/insightface) +
[HDBSCAN](https://hdbscan.readthedocs.io/) by default. To swap any
of them, point `TASKMEM_FACE_BACKEND`, `TASKMEM_FACE_CLUSTER_BACKEND`,
or (optional, for cross-clip speaker re-id) `TASKMEM_AUDIO_EMBED_BACKEND`
at your own Python module; see `tools/face_process.py` and
`tools/voice_process_lm.py` for the expected callable signatures.
## Training
RL training of the base episodic memorization policy on top of
Qwen3-VL-30B-A3B-Thinking is released as
[TaskMem-PhaseOne](https://github.com/ByteDance-Seed/TaskMem-PhaseOne). The
`training/` directory in this repo holds the data-preparation scripts
that produce the jsonl / parquet files consumed by that training
pipeline.
## Citation
```bibtex
@inproceedings{zou2026taskmem,
title = {Task-Focused Memorization for Multimodal Agents},
author = {Zou, Tao and He, Yichen and Qiu, Tian and Lin, Yuan and Li, Hang},
year = {2026},
url = {https://arxiv.org/abs/2605.31075}
}
```
## License
Apache License 2.0. See [`LICENSE`](LICENSE).