【Agentic RL专题】 深入浅出Memory模块

import faiss  
import numpy as np  
from sentence_transformers import SentenceTransformer  
  
class RAGMemoryAgent:  
    def __init__(self, embedding_dim=384, memory_size=100):  
        self.memory_size = memory_size  
        self.embedding_dim = embedding_dim  
        self.memory = []  # 外部存储的记忆  
        self.model = SentenceTransformer('paraphrase-MiniLM-L6-v2')  # 使用预训练模型  
        self.index = faiss.IndexFlatL2(self.embedding_dim)  # 向量检索索引  
  
    def encode(self, text):  
        """将文本编码为固定维度的向量（384维）"""  
        vector = self.model.encode([text])  # 生成 384 维向量  
        return vector.astype(np.float32)  
  
    def add_to_memory(self, text):  
        """将新的记忆添加到外部存储中"""  
        vector = self.encode(text)  
        self.memory.append(text)  
        self.index.add(vector)  # 将向量添加到 FAISS 索引中  
  
    def retrieve_memory(self, query, top_k=1):  
        """根据查询检索最相关的记忆"""  
        query_vector = self.encode(query)  
        D, I = self.index.search(query_vector, top_k)  # 检索最相关的记忆  
        return [self.memory[i] for i in I[0]]  # 返回检索到的记忆  
  
# 测试 RAG-style 记忆  
agent = RAGMemoryAgent(embedding_dim=384)  
agent.add_to_memory("My name is Alice and I love playing chess.")  
agent.add_to_memory("I enjoy hiking in the mountains during the weekend.")  
agent.add_to_memory("I am learning reinforcement learning to improve my skills.")  
  
query = "What is your name?"  
retrieved_memory = agent.retrieve_memory(query)  
print("Retrieved Memory:", retrieved_memory)

import numpy as np  
from sentence_transformers import SentenceTransformer  
  
class TokenMemoryAgent:  
    def __init__(self, embedding_dim=128):  
        self.memory = []  # 存储每个 token 的记忆  
        self.embedding_dim = embedding_dim  # 嵌入维度  
        self.model = SentenceTransformer('paraphrase-MiniLM-L6-v2')  # 使用预训练模型  
  
    def encode(self, text):  
        """将文本编码为嵌入向量"""  
        vector = self.model.encode([text])  # 生成向量  
        return vector[0]  # 取第一个（单一文本的嵌入）  
  
    def add_token_memory(self, token):  
        """将一个 token 存储到记忆中"""  
        token_embedding = self.encode(token)  
        self.memory.append(token_embedding)  
  
    def update_token_memory(self, token_index, new_token):  
        """更新已有的 token 记忆"""  
        token_embedding = self.encode(new_token)  
        self.memory[token_index] = token_embedding  
  
    def retrieve_memory(self, query_token):  
        """检索与查询 token 最相关的记忆"""  
        query_embedding = self.encode(query_token)  
        similarities = [np.dot(query_embedding, token) for token in self.memory]  
        best_match_index = np.argmax(similarities)  
        return self.memory[best_match_index]  
  
# 测试 Token-style 记忆  
agent = TokenMemoryAgent(embedding_dim=128)  
agent.add_token_memory("Alice loves chess.")  
agent.add_token_memory("Weekend hiking is fun.")  
  
query = "What does Alice like?"  
retrieved_memory = agent.retrieve_memory(query)  
print("Retrieved Token Memory:", retrieved_memory)

import numpy as np
from sentence_transformers import SentenceTransformer
from stable_baselines3 import PPO
from stable_baselines3.common.envs import DummyVecEnv
from stable_baselines3.common.callbacks import BaseCallback

class TokenMemoryManagementCallback(BaseCallback):
    def __init__(self, memory_agent, verbose=0):
        super().__init__(verbose)
        self.memory_agent = memory_agent

    def _on_step(self):
        # 在每一步中，使用 RL 策略来决定是否更新/删除token
        if self.training_env.get_attr("step") % 100 == 0:
            action = np.random.choice(["retain", "discard", "add"])
            if action == "add":
                self.memory_agent.add_token("New token information")
            elif action == "discard":
                self.memory_agent.remove_token()
            elif action == "retain":
                pass
        return True

class TokenMemoryAgent:
    def __init__(self, embedding_dim=128):
        self.memory = []  # 存储每个 token 的记忆
        self.embedding_dim = embedding_dim
        self.model = SentenceTransformer('paraphrase-MiniLM-L6-v2')

    def encode(self, text):
        """将文本编码为嵌入向量"""
        vector = self.model.encode([text])
        return vector[0]  # 取第一个（单一文本的嵌入）

    def add_token(self, token):
        """将一个 token 存储到记忆中"""
        token_embedding = self.encode(token)
        self.memory.append(token_embedding)

    def remove_token(self):
        """从记忆中移除一个 token"""
        if self.memory:
            self.memory.pop(np.random.choice(len(self.memory)))
            print("Token removed.")

    def retrieve_memory(self, query_token):
        """检索与查询 token 最相关的记忆"""
        query_embedding = self.encode(query_token)
        similarities = [np.dot(query_embedding, token) for token in self.memory]
        best_match_index = np.argmax(similarities)
        return self.memory[best_match_index]

# 训练环境设置
env = DummyVecEnv([lambda: gym.make('CartPole-v1')])  # 选择合适的 RL 环境
token_agent = TokenMemoryAgent()

# 使用 PPO 算法进行训练
model = PPO("MlpPolicy", env, verbose=1)
token_callback = TokenMemoryManagementCallback(token_agent)
model.learn(total_timesteps=20000, callback=token_callback)

# 测试 Token-style 记忆
query = "What should I remember?"
retrieved_memory = token_agent.retrieve_memory(query)
print("Retrieved Token Memory:", retrieved_memory)

from langchain_openai import ChatOpenAI
from langchain.chains import ConversationChain
from langchain.memory import ConversationBufferMemory

llm = ChatOpenAI(
    model="gpt-4o-mini",
    temperature=0
)

# 记录“从对话开始到现在”的完整历史
memory = ConversationBufferMemory(
    memory_key="chat_history",   # prompt 里历史的字段名
    return_messages=True         # 以 Message 对象形式返回，更灵活
)

conversation = ConversationChain(
    llm=llm,
    memory=memory,
    verbose=True  # 打印出 prompt 和中间过程，方便你写博客讲解
)

# 连续对话测试
print(conversation.invoke("你好，我叫小明，以后请记住我的名字。")["response"])
print(conversation.invoke("我刚才说我叫什么名字？")["response"])
print(conversation.invoke("再帮我用第一人称简单介绍一下自己。")["response"])

from langchain.memory import ConversationBufferWindowMemory

# 只保留最近 3 轮的对话
window_memory = ConversationBufferWindowMemory(
    k=3,
    memory_key="chat_history",
    return_messages=True
)

short_conversation = ConversationChain(
    llm=llm,
    memory=window_memory,
    verbose=True
)

# 连续问多轮，观察超过 k 轮之后早期信息会被“遗忘”
questions = [
    "1. 我叫 Alice，请记住。",
    "2. 我喜欢玩星露谷物语。",
    "3. 我讨厌夏天的蚊子。",
    "4. 你还记得我讨厌什么吗？",
    "5. 你还记得我叫什么名字吗？"
]

for q in questions:
    resp = short_conversation.invoke(q)["response"]
    print(f"Q: {q}\nA: {resp}\n{'-'*40}")