openenv/chat_env

1---2title: Chat Environment Server3emoji: 💬4colorFrom: '#0084FF'5colorTo: '#25D366'6sdk: docker7pinned: false8app_port: 80009base_path: /web10tags:11  - openenv12---13 14# Chat Environment15 16A chat-based environment for LLMs with built-in tokenization and message history management. This environment is designed to work directly with language models and provides a minimal, flexible foundation for conversation-based RL training.17 18## Overview19 20ChatEnvironment is a lightweight environment that:21- Manages conversation history in Huggingface chat format22- Handles tokenization internally using any compatible tokenizer23- Stores both messages and tokens for efficient model interaction24- Provides a clean interface for building chat-based RL agents25 26ChatEnvironment can be used in **two ways**:271. **Direct usage**: Import and use ChatEnvironment directly in your Python code (best for local development)282. **HTTP client**: Use ChatEnv client to connect to a ChatEnvironment server (best for distributed/containerized deployments)29 30## Quick Start31 32### Option 1: Direct Usage (Local)33 34```python35from transformers import AutoTokenizer36from envs.chat_env import ChatAction, ChatObservation37from envs.chat_env.server import ChatEnvironment38from openenv.core.env_server import Message39 40# Initialize with a tokenizer and optional system prompt41tokenizer = AutoTokenizer.from_pretrained("gpt2")42env = ChatEnvironment(43    tokenizer=tokenizer,44    system_prompt="You are a helpful assistant.",45    system_role="system"46)47 48# Reset the environment49obs = env.reset()50print(f"Messages: {obs.messages}")51print(f"Tokens shape: {obs.tokens.shape}")52 53# Create an action from a message54user_message: Message = {"role": "user", "content": "Hello!"}55action = env.message_to_action(user_message)56 57# Step the environment58obs = env.step(action)59print(f"Updated messages: {obs.messages}")60print(f"Updated tokens shape: {obs.tokens.shape}")61```62 63### Option 2: HTTP Client (Distributed)64 65```python66from transformers import AutoTokenizer67from envs.chat_env import ChatEnv, ChatAction68import torch69 70# Create environment from Docker image71client = ChatEnv.from_docker_image("chat-env:latest")72 73# Or connect to existing server74# client = ChatEnv(base_url="http://localhost:8000")75 76# Reset77result = client.reset()78print(f"Initial messages: {result.observation.messages}")79 80# Send an action with tokens81tokenizer = AutoTokenizer.from_pretrained("gpt2")82message = {"role": "user", "content": "Hello!"}83action = client.message_to_action(message, tokenizer)84 85result = client.step(action)86print(f"Messages: {result.observation.messages}")87print(f"Reward: {result.reward}")88 89# Cleanup90client.close()91```92 93### Building the Docker Image94 95Before using the HTTP client, build the Docker image:96 97```bash98# From project root99docker build -t chat-env:latest -f envs/chat_env/server/Dockerfile .100 101# Optionally specify a different tokenizer102docker build -t chat-env:latest \103  --build-arg TOKENIZER_NAME=meta-llama/Llama-2-7b-chat-hf \104  -f envs/chat_env/server/Dockerfile .105```106 107## Architecture108 109### Data Models110 111#### ChatAction112Actions contain only tokens (PyTorch tensors) that interface directly with models:113```python114@dataclass115class ChatAction(Action):116    tokens: torch.Tensor  # Required, cannot be empty117```118 119#### ChatObservation120Observations contain both the message history and flattened tokens:121```python122@dataclass123class ChatObservation(Observation):124    messages: list[Message]  # List of {"role": str, "content": str}125    tokens: torch.Tensor     # Flattened tensor of all conversation tokens126    # Inherited: done, reward, metadata127```128 129#### ChatState130Internal state tracking message and token history:131```python132@dataclass133class ChatState(State):134    history_messages: list[Message]135    history_tokens: list[torch.Tensor]136    # Inherited: episode_id, step_count137```138 139### Key Methods140 141#### `reset() -> ChatObservation`142Resets the environment to initial state with optional system prompt.143 144#### `step(action: ChatAction) -> ChatObservation`145Takes an action (tokens), decodes to text, adds to history, returns updated observation.146 147#### `message_to_action(message: Message) -> ChatAction`148Convenience method to convert a message dict to a tokenized ChatAction.149 150## Usage Patterns151 152### Basic Conversation153 154```python155from transformers import AutoTokenizer156from envs.chat_env.server import ChatEnvironment157from openenv.core.env_server import Message158 159tokenizer = AutoTokenizer.from_pretrained("gpt2")160env = ChatEnvironment(tokenizer=tokenizer)161 162# Reset163obs = env.reset()164 165# User turn166user_msg: Message = {"role": "user", "content": "What is 2+2?"}167action = env.message_to_action(user_msg)168obs = env.step(action)169 170# Assistant turn171assistant_msg: Message = {"role": "assistant", "content": "2+2 equals 4."}172action = env.message_to_action(assistant_msg)173obs = env.step(action)174 175# Access conversation history176print(f"Full conversation: {obs.messages}")177print(f"All tokens: {obs.tokens}")178```179 180### With Transforms181 182You can add transforms to compute rewards or modify observations:183 184```python185from openenv.core.env_server import Transform, Observation186 187class LengthRewardTransform(Transform):188    """Reward based on response length."""189 190    def __call__(self, observation: Observation) -> Observation:191        if hasattr(observation, 'messages') and observation.messages:192            last_message = observation.messages[-1]193            observation.reward = len(last_message['content']) * 0.1194        return observation195 196env = ChatEnvironment(197    tokenizer=tokenizer,198    transform=LengthRewardTransform()199)200```201 202### Direct Token Usage203 204If you're generating tokens from a model, you can create actions directly:205 206```python207import torch208from envs.chat_env import ChatAction209 210# Assume you have tokens from your model211generated_tokens = torch.tensor([[1, 2, 3, 4, 5]])212 213# Create action directly214action = ChatAction(tokens=generated_tokens)215 216# Step environment217obs = env.step(action)218```219 220## Design Philosophy221 222ChatEnvironment is intentionally minimal and flexible:223 2241. **No HTTP overhead**: Works directly with Python objects and tensors2252. **Tokenizer ownership**: Environment handles tokenization consistently2263. **Dual representation**: Maintains both human-readable messages and model-ready tokens2274. **Transform support**: Extensible reward computation and observation modification2285. **Type-safe**: Uses typed Messages compatible with Huggingface format229 230## Integration with Models231 232ChatEnvironment pairs naturally with language models:233 234```python235# Pseudo-code for RL training loop236model = YourLanguageModel()237env = ChatEnvironment(tokenizer=model.tokenizer)238 239for episode in range(num_episodes):240    obs = env.reset()241 242    while not obs.done:243        # Model generates response tokens244        action_tokens = model.generate(obs.tokens)245        action = ChatAction(tokens=action_tokens)246 247        # Step environment248        obs = env.step(action)249 250        # Use obs.reward for RL updates251        model.update(obs.reward)252```253 254## Project Structure255 256```257chat_env/258├── __init__.py              # Module exports (ChatEnv, ChatAction, etc.)259├── README.md                # This file260├── client.py                # ChatEnv HTTP client261├── models.py                # ChatAction, ChatObservation, ChatState262└── server/263    ├── __init__.py          # Server module exports264    ├── chat_environment.py  # Core ChatEnvironment implementation265    ├── app.py               # FastAPI server application266    ├── test_chat_env.py     # Unit tests267    └── Dockerfile           # Container image for HTTP server268```269 270## Requirements271 272- Python 3.10+273- PyTorch274- A tokenizer with `apply_chat_template` method (e.g., Huggingface transformers)275 276## Notes277 278- ChatEnvironment does **not** generate responses - it only manages conversation state279- You need to provide tokens from your model or other source280- The environment is thread-safe for single-threaded use only281- For multi-turn conversations, alternate between user and assistant messages282