🤖 Deploy chat_env environment - 2025-10-19 22:32:42

.gitattributes

@@ -1,35 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

Dockerfile

@@ -0,0 +1,46 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Multi-stage build: First stage builds the base image
+FROM python:3.11-slim as base-builder
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Python dependencies that all environments need
+RUN pip install --no-cache-dir \
+    fastapi>=0.104.0 \
+    "uvicorn[standard]>=0.24.0" \
+    requests>=2.25.0 \
+    wsproto>=1.0.0
+
+# Set working directory
+WORKDIR /app
+
+# Default environment variables
+ENV PYTHONPATH=/app/src
+ENV PYTHONUNBUFFERED=1
+
+# Second stage: Use the built base image and add environment-specific dependencies
+FROM base-builder
+
+# Install additional dependencies for ChatEnvironment
+RUN pip install --no-cache-dir torch transformers
+
+
+# Copy only what's needed for this environment
+COPY src/core/ /app/src/core/
+COPY src/envs/chat_env/ /app/src/envs/chat_env/
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:8000/health || exit 1
+
+# Run the FastAPI server
+CMD ["uvicorn", "envs.chat_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]
+ENV ENABLE_WEB_INTERFACE=true

README.md

@@ -1,10 +1,48 @@
 ---
-title: Chat Env
-emoji: ⚡
-colorFrom: yellow
-colorTo: blue
+title: Chat_env Environment Server
+emoji: 🐳
+colorFrom: blue
+colorTo: green
 sdk: docker
 pinned: false
+app_port: 8000
+base_path: /web
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Chat_env Environment Server
+
+FastAPI server for chat_env environment powered by Meta's OpenEnv.
+
+## About
+
+This Space provides a containerized environment for chat_env interactions.
+Built with FastAPI and OpenEnv framework.
+
+## Web Interface
+
+This deployment includes an interactive web interface for exploring the environment:
+- **HumanAgent Interface**: Interact with the environment using a web form
+- **State Observer**: Real-time view of environment state and action history
+- **Live Updates**: WebSocket-based real-time updates
+
+Access the web interface at: `/web`
+
+## Chat Environment
+
+Provides a chat-based interface for LLMs with tokenization support.
+
+### Usage
+Send a POST request to `/step` with tokenized input:
+```json
+{
+  "tokens": [1, 2, 3, 4, 5]
+}
+```
+
+## API Documentation
+
+Visit `/docs` for interactive API documentation.
+
+## Health Check
+
+The environment provides a health check endpoint at `/health`.

src/core/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Core components for agentic environments."""
+
+# Re-export main components from submodules for convenience
+from .env_server import *
+from .http_env_client import HTTPEnvClient
+from .types import StepResult
+
+# Note: MCP module doesn't export anything yet
+
+__all__ = [
+    "HTTPEnvClient",
+    "StepResult",
+]

src/core/containers/__init__.py

@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Container management for environment servers."""

src/core/containers/images/Dockerfile

@@ -0,0 +1,46 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+#
+# OpenEnv Base Image
+#
+# This is the standard base image for all OpenEnv environment servers.
+# It includes the minimal dependencies needed to run HTTP environment servers.
+#
+# Build: docker build -t openenv-base:latest -f src/core/containers/images/Dockerfile .
+# Tag:   docker tag openenv-base:latest openenv-base:0.1.0
+#
+
+FROM python:3.11-slim
+
+# Set metadata
+LABEL maintainer="OpenEnv Team"
+LABEL description="Base image for OpenEnv based environment servers"
+LABEL version="0.1.0"
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install Python dependencies that all environments need
+RUN pip install --no-cache-dir \
+    fastapi>=0.104.0 \
+    "uvicorn[standard]>=0.24.0" \
+    requests>=2.25.0 \
+    wsproto>=1.0.0
+
+# Set working directory
+WORKDIR /app
+
+# Default environment variables
+ENV PYTHONPATH=/app/src
+ENV PYTHONUNBUFFERED=1
+
+# Default expose port (can be overridden)
+EXPOSE 8000
+
+# Note: CMD should be specified in child Dockerfiles

src/core/containers/images/README.md

@@ -0,0 +1,92 @@
+# OpenEnv Base Image
+
+Standard base image for all OpenEnv environment servers.
+
+## What's Included
+
+| Layer | Size | Contents |
+|-------|------|----------|
+| python:3.11-slim | 200 MB  | Base Python runtime |
+| + Dependencies   | 100 MB  | FastAPI, uvicorn, requests |
+| **Total**        | **~300 MB** | Ready for environment servers |
+
+## Image Sizes
+
+```
+openenv-base:latest   300 MB  (python + fastapi + uvicorn)
+```
+echo-env:latest        500 MB  (python + fastapi + uvicorn + app)
+coding-env:latest      520 MB  (python + fastapi + uvicorn + app + tools)
+another-env:latest     510 MB  (python + fastapi + uvicorn + app)
+---
+Total: 1.5 GB (with lots of duplication)
+```
+
+### With Base Images (✅ Solution)
+```
+openenv-base:latest    300 MB  (python + fastapi + uvicorn)
+echo-env:latest         50 MB  (app only, uses base)
+coding-env:latest       70 MB  (app + tools, uses base)
+another-env:latest      45 MB  (app only, uses base)
+---
+Total: 465 MB (base shared, minimal duplication)
+```
+
+## Building the Base Image
+
+```bash
+# From project root
+docker build -t openenv-base:latest -f src/core/containers/images/Dockerfile .
+```
+
+## Usage in Environment Dockerfiles
+
+Each environment Dockerfile should start with:
+
+```dockerfile
+FROM openenv-base:latest
+
+# Copy only environment-specific files
+COPY src/core/ /app/src/core/
+COPY src/envs/my_env/ /app/src/envs/my_env/
+
+# Run the server
+CMD ["uvicorn", "envs.my_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]
+```
+
+## Base Image Contents
+
+- Python 3.11-slim
+- FastAPI >= 0.104.0
+- Uvicorn >= 0.24.0
+- Requests >= 2.25.0
+- curl (for health checks)
+
+## Example: Building Echo Environment
+
+```bash
+# Step 1: Build base image (do this once)
+docker build -t openenv-base:latest -f src/core/containers/images/Dockerfile .
+
+# Step 2: Build echo environment (uses base)
+docker build -t echo-env:latest -f src/envs/echo_env/server/Dockerfile .
+
+# Step 3: Run echo environment
+docker run -p 8000:8000 echo-env:latest
+```
+
+## Updating the Base
+
+When dependencies need updating:
+
+1. Update `src/core/containers/images/Dockerfile`
+2. Rebuild base image
+3. Rebuild all environment images (they'll use new base)
+
+```bash
+# Update base
+docker build -t openenv-base:latest -f src/core/containers/images/Dockerfile .
+
+# Rebuild environments (they automatically use new base)
+docker build -t echo-env:latest -f src/envs/echo_env/server/Dockerfile .
+```

src/core/containers/runtime/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Container runtime providers."""
+
+from .providers import ContainerProvider, KubernetesProvider, LocalDockerProvider
+
+__all__ = [
+    "ContainerProvider",
+    "LocalDockerProvider",
+    "KubernetesProvider",
+]

src/core/containers/runtime/providers.py

@@ -0,0 +1,289 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Container provider abstractions for running environment servers.
+
+This module provides a pluggable architecture for different container providers
+(local Docker, Kubernetes, cloud providers, etc.) to be used with HTTPEnvClient.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional
+
+
+class ContainerProvider(ABC):
+    """
+    Abstract base class for container providers.
+
+    Providers implement this interface to support different container platforms:
+    - LocalDockerProvider: Runs containers on local Docker daemon
+    - KubernetesProvider: Runs containers in Kubernetes cluster
+    - FargateProvider: Runs containers on AWS Fargate
+    - CloudRunProvider: Runs containers on Google Cloud Run
+
+    The provider manages a single container lifecycle and provides the base URL
+    for connecting to it.
+
+    Example:
+        >>> provider = LocalDockerProvider()
+        >>> base_url = provider.start_container("echo-env:latest")
+        >>> print(base_url)  # http://localhost:8000
+        >>> # Use the environment via base_url
+        >>> provider.stop_container()
+    """
+
+    @abstractmethod
+    def start_container(
+        self,
+        image: str,
+        port: Optional[int] = None,
+        env_vars: Optional[Dict[str, str]] = None,
+        **kwargs: Any,
+    ) -> str:
+        """
+        Start a container from the specified image.
+
+        Args:
+            image: Container image name (e.g., "echo-env:latest")
+            port: Port to expose (if None, provider chooses)
+            env_vars: Environment variables to pass to container
+            **kwargs: Provider-specific options
+
+        Returns:
+            Base URL to connect to the container (e.g., "http://localhost:8000")
+
+        Raises:
+            RuntimeError: If container fails to start
+        """
+        pass
+
+    @abstractmethod
+    def stop_container(self) -> None:
+        """
+        Stop and remove the running container.
+
+        This cleans up the container that was started by start_container().
+        """
+        pass
+
+    @abstractmethod
+    def wait_for_ready(self, base_url: str, timeout_s: float = 30.0) -> None:
+        """
+        Wait for the container to be ready to accept requests.
+
+        This typically polls the /health endpoint until it returns 200.
+
+        Args:
+            base_url: Base URL of the container
+            timeout_s: Maximum time to wait
+
+        Raises:
+            TimeoutError: If container doesn't become ready in time
+        """
+        pass
+
+
+class LocalDockerProvider(ContainerProvider):
+    """
+    Container provider for local Docker daemon.
+
+    This provider runs containers on the local machine using Docker.
+    Useful for development and testing.
+
+    Example:
+        >>> provider = LocalDockerProvider()
+        >>> base_url = provider.start_container("echo-env:latest")
+        >>> # Container running on http://localhost:<random-port>
+        >>> provider.stop_container()
+    """
+
+    def __init__(self):
+        """Initialize the local Docker provider."""
+        self._container_id: Optional[str] = None
+        self._container_name: Optional[str] = None
+
+        # Check if Docker is available
+        import subprocess
+
+        try:
+            subprocess.run(
+                ["docker", "version"],
+                check=True,
+                capture_output=True,
+                timeout=5,
+            )
+        except (subprocess.CalledProcessError, FileNotFoundError, subprocess.TimeoutExpired):
+            raise RuntimeError(
+                "Docker is not available. Please install Docker Desktop or Docker Engine."
+            )
+
+    def start_container(
+        self,
+        image: str,
+        port: Optional[int] = None,
+        env_vars: Optional[Dict[str, str]] = None,
+        **kwargs: Any,
+    ) -> str:
+        """
+        Start a Docker container locally.
+
+        Args:
+            image: Docker image name
+            port: Port to expose (if None, finds available port)
+            env_vars: Environment variables for the container
+            **kwargs: Additional Docker run options
+
+        Returns:
+            Base URL to connect to the container
+        """
+        import subprocess
+        import time
+
+        # Find available port if not specified
+        if port is None:
+            port = self._find_available_port()
+
+        # Generate container name
+        self._container_name = self._generate_container_name(image)
+
+        # Build docker run command
+        cmd = [
+            "docker", "run",
+            "-d",  # Detached
+            "--name", self._container_name,
+            "-p", f"{port}:8000",  # Map port
+        ]
+
+        # Add environment variables
+        if env_vars:
+            for key, value in env_vars.items():
+                cmd.extend(["-e", f"{key}={value}"])
+
+        # Add image
+        cmd.append(image)
+
+        # Run container
+        result = subprocess.run(cmd, capture_output=True, text=True, check=True)
+        self._container_id = result.stdout.strip()
+
+        # Wait a moment for container to start
+        time.sleep(1)
+
+        base_url = f"http://localhost:{port}"
+        return base_url
+
+    def stop_container(self) -> None:
+        """
+        Stop and remove the Docker container.
+        """
+        if self._container_id is None:
+            return
+
+        import subprocess
+
+        try:
+            # Stop container
+            subprocess.run(
+                ["docker", "stop", self._container_id],
+                capture_output=True,
+                check=True,
+                timeout=10,
+            )
+
+            # Remove container
+            subprocess.run(
+                ["docker", "rm", self._container_id],
+                capture_output=True,
+                check=True,
+                timeout=10,
+            )
+        except subprocess.CalledProcessError:
+            # Container might already be stopped/removed
+            pass
+        finally:
+            self._container_id = None
+            self._container_name = None
+
+    def wait_for_ready(self, base_url: str, timeout_s: float = 30.0) -> None:
+        """
+        Wait for container to be ready by polling /health endpoint.
+
+        Args:
+            base_url: Base URL of the container
+            timeout_s: Maximum time to wait
+
+        Raises:
+            TimeoutError: If container doesn't become ready
+        """
+        import time
+        import requests
+
+        start_time = time.time()
+        health_url = f"{base_url}/health"
+
+        while time.time() - start_time < timeout_s:
+            try:
+                response = requests.get(health_url, timeout=2.0)
+                if response.status_code == 200:
+                    return
+            except requests.RequestException:
+                pass
+
+            time.sleep(0.5)
+
+        raise TimeoutError(
+            f"Container at {base_url} did not become ready within {timeout_s}s"
+        )
+
+    def _find_available_port(self) -> int:
+        """
+        Find an available port on localhost.
+
+        Returns:
+            An available port number
+        """
+        import socket
+
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(("", 0))
+            s.listen(1)
+            port = s.getsockname()[1]
+        return port
+
+    def _generate_container_name(self, image: str) -> str:
+        """
+        Generate a unique container name based on image name and timestamp.
+
+        Args:
+            image: Docker image name
+
+        Returns:
+            A unique container name
+        """
+        import time
+
+        clean_image = image.split("/")[-1].split(":")[0]
+        timestamp = int(time.time() * 1000)
+        return f"{clean_image}-{timestamp}"
+
+
+class KubernetesProvider(ContainerProvider):
+    """
+    Container provider for Kubernetes clusters.
+
+    This provider creates pods in a Kubernetes cluster and exposes them
+    via services or port-forwarding.
+
+    Example:
+        >>> provider = KubernetesProvider(namespace="envtorch-dev")
+        >>> base_url = provider.start_container("echo-env:latest")
+        >>> # Pod running in k8s, accessible via service or port-forward
+        >>> provider.stop_container()
+    """
+    pass

src/core/containers/test_local_docker_provider.py

@@ -0,0 +1,258 @@
+#!/usr/bin/env python3
+"""
+End-to-end test for LocalDockerProvider.
+
+This script tests the complete flow:
+1. Start a container using LocalDockerProvider
+2. Wait for it to be ready
+3. Make HTTP requests to test the environment
+4. Clean up the container
+"""
+
+import sys
+from pathlib import Path
+
+# Add src to path
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+import requests
+
+from core.containers.runtime import LocalDockerProvider
+
+# TODO: Remove this test or make it a functional test sicne this will be tested in e2e test for echo env
+def test_local_docker_provider():
+    """Test LocalDockerProvider end-to-end."""
+    print("=" * 60)
+    print("LocalDockerProvider End-to-End Test")
+    print("=" * 60)
+    print()
+
+    provider = None
+
+    try:
+        # Step 1: Create provider
+        print("Step 1: Creating LocalDockerProvider...")
+        provider = LocalDockerProvider()
+        print("✓ Provider created\n")
+
+        # Step 2: Start container
+        print("Step 2: Starting echo-env container...")
+        base_url = provider.start_container("echo-env:latest")
+        print(f"✓ Container started at: {base_url}")
+        if provider._container_id:
+            print(f"  Container ID: {provider._container_id[:12]}...")
+        if provider._container_name:
+            print(f"  Container name: {provider._container_name}\n")
+
+        # Step 3: Wait for ready
+        print("Step 3: Waiting for container to be ready...")
+        provider.wait_for_ready(base_url, timeout_s=30.0)
+        print("✓ Container is ready!\n")
+
+        # Step 4: Test health endpoint
+        print("Step 4: Testing /health endpoint...")
+        response = requests.get(f"{base_url}/health")
+        print(f"  Status: {response.status_code}")
+        print(f"  Response: {response.json()}")
+        assert response.status_code == 200
+        assert response.json()["status"] == "healthy"
+        print("✓ Health check passed\n")
+
+        # Step 5: Test reset endpoint
+        print("Step 5: Testing /reset endpoint...")
+        response = requests.post(
+            f"{base_url}/reset",
+            json={},
+            headers={"Content-Type": "application/json"},
+        )
+        print(f"  Status: {response.status_code}")
+        data = response.json()
+        print(f"  Message: {data['observation']['echoed_message']}")
+        print(f"  Reward: {data['reward']}")
+        print(f"  Done: {data['done']}")
+        assert response.status_code == 200
+        assert data["observation"]["echoed_message"] == "Echo environment ready!"
+        print("✓ Reset test passed\n")
+
+        # Step 6: Test step endpoint
+        print("Step 6: Testing /step endpoint...")
+        response = requests.post(
+            f"{base_url}/step",
+            json={"action": {"message": "Hello from LocalDockerProvider!"}},
+            headers={"Content-Type": "application/json"},
+        )
+        print(f"  Status: {response.status_code}")
+        data = response.json()
+        print(f"  Echoed: {data['observation']['echoed_message']}")
+        print(f"  Length: {data['observation']['message_length']}")
+        print(f"  Reward: {data['reward']}")
+        assert response.status_code == 200
+        assert data["observation"]["echoed_message"] == "Hello from LocalDockerProvider!"
+        assert data["observation"]["message_length"] == 31
+        print("✓ Step test passed\n")
+
+        # Step 7: Test state endpoint
+        print("Step 7: Testing /state endpoint...")
+        response = requests.get(f"{base_url}/state")
+        print(f"  Status: {response.status_code}")
+        data = response.json()
+        print(f"  Episode ID: {data['episode_id']}")
+        print(f"  Step count: {data['step_count']}")
+        assert response.status_code == 200
+        assert data["step_count"] == 1  # One step from above
+        print("✓ State test passed\n")
+
+        # Step 8: Multiple steps
+        print("Step 8: Testing multiple steps...")
+        for i in range(3):
+            response = requests.post(
+                f"{base_url}/step",
+                json={"action": {"message": f"Message {i+1}"}},
+                headers={"Content-Type": "application/json"},
+            )
+            assert response.status_code == 200
+            print(f"  Step {i+1}: ✓")
+
+        # Check state updated
+        response = requests.get(f"{base_url}/state")
+        data = response.json()
+        assert data["step_count"] == 4  # 1 + 3 more steps
+        print(f"  Final step count: {data['step_count']}")
+        print("✓ Multiple steps test passed\n")
+
+        print("=" * 60)
+        print("✓ All tests passed!")
+        print("=" * 60)
+        print()
+
+        return True
+
+    except Exception as e:
+        print(f"\n❌ Test failed: {e}")
+        import traceback
+        traceback.print_exc()
+        return False
+
+    finally:
+        # Step 9: Cleanup
+        if provider is not None:
+            print("\nStep 9: Cleaning up container...")
+            try:
+                provider.stop_container()
+                print("✓ Container stopped and removed\n")
+            except Exception as e:
+                print(f"⚠️  Cleanup warning: {e}\n")
+
+
+def test_provider_with_custom_port():
+    """Test provider with custom port."""
+    print("=" * 60)
+    print("LocalDockerProvider with Custom Port Test")
+    print("=" * 60)
+    print()
+
+    provider = None
+
+    try:
+        provider = LocalDockerProvider()
+
+        print("Starting container on custom port 8123...")
+        base_url = provider.start_container("echo-env:latest", port=8123)
+        print(f"✓ Started at: {base_url}")
+        assert ":8123" in base_url
+
+        print("Waiting for ready...")
+        provider.wait_for_ready(base_url)
+        print("✓ Ready!")
+
+        print("Testing health...")
+        response = requests.get(f"{base_url}/health")
+        assert response.status_code == 200
+        print("✓ Health check passed")
+
+        print("\n✓ Custom port test passed!\n")
+        return True
+
+    except Exception as e:
+        print(f"\n❌ Test failed: {e}")
+        return False
+
+    finally:
+        if provider is not None:
+            provider.stop_container()
+            print("✓ Cleaned up\n")
+
+
+def test_provider_with_env_vars():
+    """Test provider with environment variables."""
+    print("=" * 60)
+    print("LocalDockerProvider with Environment Variables Test")
+    print("=" * 60)
+    print()
+
+    provider = None
+
+    try:
+        provider = LocalDockerProvider()
+
+        print("Starting container with environment variables...")
+        base_url = provider.start_container(
+            "echo-env:latest",
+            env_vars={"DEBUG": "true", "LOG_LEVEL": "info"}
+        )
+        print(f"✓ Started at: {base_url}")
+
+        print("Waiting for ready...")
+        provider.wait_for_ready(base_url)
+        print("✓ Ready!")
+
+        print("Testing health...")
+        response = requests.get(f"{base_url}/health")
+        assert response.status_code == 200
+        print("✓ Health check passed")
+
+        print("\n✓ Environment variables test passed!\n")
+        return True
+
+    except Exception as e:
+        print(f"\n❌ Test failed: {e}")
+        return False
+
+    finally:
+        if provider is not None:
+            provider.stop_container()
+            print("✓ Cleaned up\n")
+
+
+if __name__ == "__main__":
+    print()
+    print("🐳 LocalDockerProvider Test Suite")
+    print()
+
+    results = []
+
+    # Run basic test
+    results.append(("Basic End-to-End", test_local_docker_provider()))
+
+    # Run custom port test
+    results.append(("Custom Port", test_provider_with_custom_port()))
+
+    # Run environment variables test
+    results.append(("Environment Variables", test_provider_with_env_vars()))
+
+    # Summary
+    print("=" * 60)
+    print("Test Summary")
+    print("=" * 60)
+    for name, passed in results:
+        status = "✓ PASSED" if passed else "✗ FAILED"
+        print(f"{name:25} {status}")
+    print("=" * 60)
+
+    all_passed = all(result for _, result in results)
+    if all_passed:
+        print("\n🎉 All tests passed!")
+        exit(0)
+    else:
+        print("\n❌ Some tests failed")
+        exit(1)

src/core/env_server/__init__.py

@@ -0,0 +1,35 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Core environment interfaces and types."""
+
+from .base_transforms import CompositeTransform, NullTransform
+from .http_server import HTTPEnvServer, create_app, create_fastapi_app
+from .interfaces import Environment, Message, ModelTokenizer, Transform
+from .types import Action, Observation, State
+from .web_interface import create_web_interface_app, WebInterfaceManager
+
+__all__ = [
+    # Core interfaces
+    "Environment",
+    "Transform",
+    "Message",
+    "ModelTokenizer",
+    # Types
+    "Action",
+    "Observation",
+    "State",
+    # Base transforms
+    "CompositeTransform",
+    "NullTransform",
+    # HTTP Server
+    "HTTPEnvServer",
+    "create_app",
+    "create_fastapi_app",
+    # Web Interface
+    "create_web_interface_app",
+    "WebInterfaceManager",
+]

src/core/env_server/base_transforms.py

@@ -0,0 +1,29 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Base transform implementations for composing environment-specific transforms."""
+
+from .interfaces import Transform
+from .types import Observation
+
+
+class CompositeTransform(Transform):
+    """Combines multiple transforms into a single transform."""
+
+    def __init__(self, transforms: list[Transform]):
+        self.transforms = transforms
+
+    def __call__(self, observation: Observation) -> Observation:
+        for transform in self.transforms:
+            observation = transform(observation)
+        return observation
+
+
+class NullTransform(Transform):
+    """Default transform that passes through unchanged."""
+
+    def __call__(self, observation: Observation) -> Observation:
+        return observation

src/core/env_server/http_server.py

@@ -0,0 +1,231 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+HTTP server wrapper for Environment instances.
+
+This module provides utilities to wrap any Environment subclass and expose it
+over HTTP endpoints that HTTPEnvClient can consume.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import asdict
+from typing import Any, Dict, Type
+
+from .interfaces import Environment
+from .types import Action, Observation
+from fastapi import Body, FastAPI
+
+class HTTPEnvServer:
+    """
+    HTTP server wrapper for Environment instances.
+
+    This class wraps an Environment and exposes its reset(), step(), and state
+    methods as HTTP endpoints compatible with HTTPEnvClient.
+
+    The server expects:
+    - Action deserialization: Converts JSON dict to Action subclass
+    - Observation serialization: Converts Observation subclass to JSON dict
+
+    Example:
+        >>> from core.env_server import HTTPEnvServer
+        >>> from envs.coding_env.server import CodeExecutionEnvironment
+        >>>
+        >>> env = CodeExecutionEnvironment()
+        >>> server = HTTPEnvServer(env)
+        >>>
+        >>> # Register routes with FastAPI
+        >>> from fastapi import FastAPI
+        >>> app = FastAPI()
+        >>> server.register_routes(app)
+    """
+
+    def __init__(
+        self,
+        env: Environment,
+        action_cls: Type[Action],
+        observation_cls: Type[Observation],
+    ):
+        """
+        Initialize HTTP server wrapper.
+
+        Args:
+            env: The Environment instance to wrap
+            action_cls: The Action subclass this environment expects
+            observation_cls: The Observation subclass this environment returns
+        """
+        self.env = env
+        self.action_cls = action_cls
+        self.observation_cls = observation_cls
+
+    def register_routes(self, app: Any) -> None:
+        """
+        Register HTTP routes on a FastAPI application.
+
+        Args:
+            app: FastAPI application instance
+        """
+
+        if not isinstance(app, FastAPI):
+            raise TypeError("app must be a FastAPI instance")
+
+        @app.post("/reset")
+        async def reset(request: Dict[str, Any] = Body(default={})) -> Dict[str, Any]:
+            """Reset endpoint - returns initial observation."""
+            # TODO: Handle seed, episode_id from request if provided
+            observation = self.env.reset()
+            return self._serialize_observation(observation)
+
+        @app.post("/step")
+        async def step(request: Dict[str, Any]) -> Dict[str, Any]:
+            """Step endpoint - executes action and returns observation."""
+            action_data = request.get("action", {})
+            # TODO: Handle timeout_s, request_id, episode_id from request if provided
+
+            # Deserialize action
+            action = self._deserialize_action(action_data)
+
+            # Execute step
+            observation = self.env.step(action)
+
+            # Return serialized observation
+            return self._serialize_observation(observation)
+
+        @app.get("/state")
+        async def get_state() -> Dict[str, Any]:
+            """State endpoint - returns current environment state."""
+            state = self.env.state
+            return asdict(state)
+
+        @app.get("/health")
+        async def health() -> Dict[str, str]:
+            """Health check endpoint."""
+            return {"status": "healthy"}
+
+
+    def _deserialize_action(self, action_data: Dict[str, Any]) -> Action:
+        """
+        Convert JSON dict to Action instance.
+
+        Args:
+            action_data: Dictionary containing action data
+
+        Returns:
+            Action instance
+
+        Note:
+            This is a simple implementation. Subclasses may need to override
+            for more complex deserialization logic.
+        """
+        # Remove metadata if present (it will be set via kw_only field)
+        metadata = action_data.pop("metadata", {})
+        action = self.action_cls(**action_data)
+        action.metadata = metadata
+        return action
+
+    def _serialize_observation(self, observation: Observation) -> Dict[str, Any]:
+        """
+        Convert Observation instance to JSON-compatible dict.
+
+        Args:
+            observation: Observation instance
+
+        Returns:
+            Dictionary compatible with HTTPEnvClient._parse_result()
+
+        The format matches what HTTPEnvClient expects:
+        {
+            "observation": {...},  # Observation fields
+            "reward": float | None,
+            "done": bool,
+        }
+        """
+        obs_dict = asdict(observation)
+
+        # Extract reward and done (these are part of StepResult on client side)
+        reward = obs_dict.pop("reward", None)
+        done = obs_dict.pop("done", False)
+        obs_dict.pop("metadata", None)  # Remove metadata from observation
+
+        # Return in HTTPEnvClient expected format
+        return {
+            "observation": obs_dict,
+            "reward": reward,
+            "done": done,
+        }
+
+def create_app(
+    env: Environment,
+    action_cls: Type[Action],
+    observation_cls: Type[Observation],
+) -> Any:
+    """
+    Create a FastAPI application with web interface enabled for Hugging Face deployments.
+    
+    This function checks for the ENABLE_WEB_INTERFACE environment variable to determine
+    whether to enable the web interface. 
+    
+    Args:
+        env: The Environment instance to serve
+        action_cls: The Action subclass this environment expects
+        observation_cls: The Observation subclass this environment returns
+        
+    Returns:
+        FastAPI application instance with or without web interface based on environment
+    """
+    # Check if web interface should be enabled
+    # This can be controlled via environment variable or build argument
+    enable_web = (
+        os.getenv("ENABLE_WEB_INTERFACE", "false").lower() in ("true", "1", "yes")
+    )
+
+    if enable_web:
+        # Import web interface only when needed
+        from .web_interface import create_web_interface_app
+        return create_web_interface_app(env, action_cls, observation_cls)
+    else:
+        # Use standard FastAPI app without web interface
+        return create_fastapi_app(env, action_cls, observation_cls)
+    
+
+def create_fastapi_app(
+    env: Environment,
+    action_cls: Type[Action],
+    observation_cls: Type[Observation],
+) -> Any:
+    """
+    Create a FastAPI application with routes for the given environment.
+
+    Args:
+        env: The Environment instance to serve
+        action_cls: The Action subclass this environment expects
+        observation_cls: The Observation subclass this environment returns
+
+    Returns:
+        FastAPI application instance with routes registered
+
+    Example:
+        >>> from envs.coding_env.server import CodeExecutionEnvironment
+        >>> from envs.coding_env.models import CodeAction, CodeObservation
+        >>>
+        >>> env = CodeExecutionEnvironment()
+        >>> app = create_fastapi_app(env, CodeAction, CodeObservation)
+        >>>
+        >>> # Run with: uvicorn module:app --host 0.0.0.0 --port 8000
+    """
+    try:
+        from fastapi import FastAPI
+    except ImportError:
+        raise ImportError(
+            "FastAPI is required. Install with: pip install fastapi uvicorn"
+        )
+
+    app = FastAPI(title="Environment HTTP Server")
+    server = HTTPEnvServer(env, action_cls, observation_cls)
+    server.register_routes(app)
+    return app

src/core/env_server/interfaces.py

@@ -0,0 +1,118 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from abc import ABC, abstractmethod
+from typing import Any, Protocol, TypedDict
+
+from .types import Action, Observation, State
+
+
+class Message(TypedDict):
+    """A message in a conversation.
+
+    Compatible with Huggingface chat template format.
+    """
+
+    role: str
+    content: str
+
+
+class ModelTokenizer(Protocol):
+    """Protocol for tokenizers that support chat templates.
+
+    This protocol defines the interface that tokenizers must implement
+    to work with chat-based environments. It's compatible with
+    Huggingface transformers tokenizers.
+    """
+
+    def apply_chat_template(
+        self,
+        conversation: list[Message],
+        tokenize: bool = True,
+        return_tensors: str | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Apply a chat template to format and optionally tokenize a conversation.
+
+        Args:
+            conversation: List of message dictionaries with 'role' and 'content'
+            tokenize: Whether to tokenize the output
+            return_tensors: Format for returned tensors ('pt' for PyTorch)
+            **kwargs: Additional arguments
+
+        Returns:
+            Formatted and optionally tokenized conversation
+        """
+        ...
+
+    def decode(
+        self, token_ids: Any, skip_special_tokens: bool = False, **kwargs: Any
+    ) -> str:
+        """Decode token IDs back to text.
+
+        Args:
+            token_ids: Token IDs to decode
+            skip_special_tokens: Whether to skip special tokens in output
+            **kwargs: Additional arguments
+
+        Returns:
+            Decoded text string
+        """
+        ...
+
+
+class Transform(ABC):
+    """Transform observations to add rewards, metrics, or other modifications.
+
+    Transforms follow the TorchRL pattern where they take an observation
+    and return a (potentially modified) observation. This allows for
+    flexible reward computation and observation augmentation.
+    """
+
+    @abstractmethod
+    def __call__(self, observation: Observation) -> Observation:
+        """Transform an observation.
+
+        Args:
+            observation: The input observation
+
+        Returns:
+            The transformed observation
+        """
+        pass
+
+
+class Environment(ABC):
+    """Base class for all environment servers following Gym/Gymnasium API.
+
+    Args:
+        transform: Optional transform to apply to observations
+    """
+
+    def __init__(self, transform: Transform | None = None):
+        self.transform = transform
+
+    @abstractmethod
+    def reset(self) -> Observation:
+        """Reset the environment and return initial observation."""
+        pass
+
+    @abstractmethod
+    def step(self, action: Action) -> Observation:
+        """Take a step in the environment."""
+        pass
+
+    @property
+    @abstractmethod
+    def state(self) -> State:
+        """Get the current environment state."""
+        pass
+
+    def _apply_transform(self, observation: Observation) -> Observation:
+        """Apply transform if one is provided."""
+        if self.transform is not None:
+            return self.transform(observation)
+        return observation

src/core/env_server/types.py

@@ -0,0 +1,45 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Union
+
+
+# Type aliases
+Scalar = Union[int, float, bool]
+
+
+@dataclass(kw_only=True)
+class Action:
+    """Base class for all environment actions."""
+
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass(kw_only=True)
+class Observation:
+    """Base class for all environment observations."""
+
+    done: bool = False
+    reward: Union[bool, int, float, None] = None
+    metadata: Dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class State:
+    """Base class for environment state."""
+
+    episode_id: Optional[str] = None
+    step_count: int = 0
+
+
+@dataclass
+class CodeExecResult:
+    """Result of code execution containing stdout, stderr, and exit code."""
+
+    stdout: str
+    stderr: str
+    exit_code: int

src/core/env_server/web_interface.py

@@ -0,0 +1,764 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Web interface for OpenEnv environments.
+
+This module provides a web-based interface for interacting with OpenEnv environments,
+including a two-pane layout for HumanAgent interaction and state observation.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional, Type
+from datetime import datetime
+
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
+from fastapi.responses import HTMLResponse, FileResponse
+from fastapi.staticfiles import StaticFiles
+from pydantic import BaseModel
+
+from .interfaces import Environment
+from .types import Action, Observation, State
+
+
+@dataclass
+class ActionLog:
+    """Log entry for an action taken."""
+    timestamp: str
+    action: Dict[str, Any]
+    observation: Dict[str, Any]
+    reward: Optional[float]
+    done: bool
+    step_count: int
+
+
+@dataclass
+class EpisodeState:
+    """Current episode state for the web interface."""
+    episode_id: Optional[str]
+    step_count: int
+    current_observation: Optional[Dict[str, Any]]
+    action_logs: List[ActionLog]
+    is_reset: bool = True
+
+
+class WebInterfaceManager:
+    """Manages the web interface for an environment."""
+    
+    def __init__(
+        self,
+        env: Environment,
+        action_cls: Type[Action],
+        observation_cls: Type[Observation],
+    ):
+        self.env = env
+        self.action_cls = action_cls
+        self.observation_cls = observation_cls
+        self.episode_state = EpisodeState(
+            episode_id=None,
+            step_count=0,
+            current_observation=None,
+            action_logs=[]
+        )
+        self.connected_clients: List[WebSocket] = []
+    
+    async def connect_websocket(self, websocket: WebSocket):
+        """Connect a new WebSocket client."""
+        await websocket.accept()
+        self.connected_clients.append(websocket)
+        
+        # Send current state to the new client
+        await self._send_state_update()
+    
+    async def disconnect_websocket(self, websocket: WebSocket):
+        """Disconnect a WebSocket client."""
+        if websocket in self.connected_clients:
+            self.connected_clients.remove(websocket)
+    
+    async def _send_state_update(self):
+        """Send current state to all connected clients."""
+        if not self.connected_clients:
+            return
+            
+        state_data = {
+            "type": "state_update",
+            "episode_state": asdict(self.episode_state)
+        }
+        
+        # Send to all connected clients
+        disconnected_clients = []
+        for client in self.connected_clients:
+            try:
+                await client.send_text(json.dumps(state_data))
+            except:
+                disconnected_clients.append(client)
+        
+        # Remove disconnected clients
+        for client in disconnected_clients:
+            self.connected_clients.remove(client)
+    
+    async def reset_environment(self) -> Dict[str, Any]:
+        """Reset the environment and update state."""
+        observation = self.env.reset()
+        state = self.env.state
+        
+        # Update episode state
+        self.episode_state.episode_id = state.episode_id
+        self.episode_state.step_count = 0
+        self.episode_state.current_observation = asdict(observation)
+        self.episode_state.action_logs = []
+        self.episode_state.is_reset = True
+        
+        # Send state update
+        await self._send_state_update()
+        
+        return {
+            "observation": asdict(observation),
+            "reward": observation.reward,
+            "done": observation.done,
+        }
+    
+    async def step_environment(self, action_data: Dict[str, Any]) -> Dict[str, Any]:
+        """Execute a step in the environment and update state."""
+        # Deserialize action
+        action = self._deserialize_action(action_data)
+        
+        # Execute step
+        observation = self.env.step(action)
+        state = self.env.state
+        
+        # Create action log
+        action_log = ActionLog(
+            timestamp=datetime.now().isoformat(),
+            action=asdict(action),
+            observation=asdict(observation),
+            reward=observation.reward,
+            done=observation.done,
+            step_count=state.step_count
+        )
+        
+        # Update episode state
+        self.episode_state.episode_id = state.episode_id
+        self.episode_state.step_count = state.step_count
+        self.episode_state.current_observation = asdict(observation)
+        self.episode_state.action_logs.append(action_log)
+        self.episode_state.is_reset = False
+        
+        # Send state update
+        await self._send_state_update()
+        
+        return {
+            "observation": asdict(observation),
+            "reward": observation.reward,
+            "done": observation.done,
+        }
+    
+    def get_state(self) -> Dict[str, Any]:
+        """Get current environment state."""
+        state = self.env.state
+        return asdict(state)
+    
+    def _deserialize_action(self, action_data: Dict[str, Any]) -> Action:
+        """Convert JSON dict to Action instance."""
+        metadata = action_data.pop("metadata", {})
+        action = self.action_cls(**action_data)
+        action.metadata = metadata
+        return action
+
+
+def create_web_interface_app(
+    env: Environment,
+    action_cls: Type[Action],
+    observation_cls: Type[Observation],
+) -> FastAPI:
+    """
+    Create a FastAPI application with web interface for the given environment.
+    
+    Args:
+        env: The Environment instance to serve
+        action_cls: The Action subclass this environment expects
+        observation_cls: The Observation subclass this environment returns
+        
+    Returns:
+        FastAPI application instance with web interface
+    """
+    from .http_server import create_fastapi_app
+    
+    # Create the base environment app
+    app = create_fastapi_app(env, action_cls, observation_cls)
+    
+    # Create web interface manager
+    web_manager = WebInterfaceManager(env, action_cls, observation_cls)
+    
+    # Add web interface routes
+    @app.get("/web", response_class=HTMLResponse)
+    async def web_interface():
+        """Serve the web interface."""
+        return get_web_interface_html(action_cls)
+    
+    @app.websocket("/ws")
+    async def websocket_endpoint(websocket: WebSocket):
+        """WebSocket endpoint for real-time updates."""
+        await web_manager.connect_websocket(websocket)
+        try:
+            while True:
+                # Keep connection alive
+                await websocket.receive_text()
+        except WebSocketDisconnect:
+            await web_manager.disconnect_websocket(websocket)
+    
+    @app.post("/web/reset")
+    async def web_reset():
+        """Reset endpoint for web interface."""
+        return await web_manager.reset_environment()
+    
+    @app.post("/web/step")
+    async def web_step(request: Dict[str, Any]):
+        """Step endpoint for web interface."""
+        action_data = request.get("action", {})
+        return await web_manager.step_environment(action_data)
+    
+    @app.get("/web/state")
+    async def web_state():
+        """State endpoint for web interface."""
+        return web_manager.get_state()
+    
+    return app
+
+
+def get_web_interface_html(action_cls: Type[Action]) -> str:
+    """Generate the HTML for the web interface."""
+    
+    # Get action fields for dynamic form generation
+    action_fields = []
+    if hasattr(action_cls, '__dataclass_fields__'):
+        for field_name, field_info in action_cls.__dataclass_fields__.items():
+            if field_name != 'metadata':
+                field_type = field_info.type
+                if field_type == str:
+                    input_type = "text"
+                elif field_type == int:
+                    input_type = "number"
+                elif field_type == float:
+                    input_type = "number"
+                elif field_type == bool:
+                    input_type = "checkbox"
+                else:
+                    input_type = "text"
+                
+                action_fields.append({
+                    'name': field_name,
+                    'type': input_type,
+                    'required': field_info.default is field_info.default_factory
+                })
+    
+    return f"""
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>OpenEnv Web Interface</title>
+    <style>
+        * {{
+            margin: 0;
+            padding: 0;
+            box-sizing: border-box;
+        }}
+        
+        body {{
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+            background-color: #f5f5f5;
+            height: 100vh;
+            overflow: hidden;
+        }}
+        
+        .container {{
+            display: flex;
+            height: 100vh;
+        }}
+        
+        .left-pane {{
+            width: 50%;
+            background: white;
+            border-right: 1px solid #e0e0e0;
+            display: flex;
+            flex-direction: column;
+        }}
+        
+        .right-pane {{
+            width: 50%;
+            background: #fafafa;
+            display: flex;
+            flex-direction: column;
+        }}
+        
+        .pane-header {{
+            padding: 20px;
+            border-bottom: 1px solid #e0e0e0;
+            background: #f8f9fa;
+            font-weight: 600;
+            font-size: 16px;
+        }}
+        
+        .pane-content {{
+            flex: 1;
+            padding: 20px;
+            overflow-y: auto;
+        }}
+        
+        .action-form {{
+            background: white;
+            border: 1px solid #e0e0e0;
+            border-radius: 8px;
+            padding: 20px;
+            margin-bottom: 20px;
+        }}
+        
+        .form-group {{
+            margin-bottom: 15px;
+        }}
+        
+        .form-group label {{
+            display: block;
+            margin-bottom: 5px;
+            font-weight: 500;
+            color: #333;
+        }}
+        
+        .form-group input, .form-group textarea {{
+            width: 100%;
+            padding: 8px 12px;
+            border: 1px solid #ddd;
+            border-radius: 4px;
+            font-size: 14px;
+        }}
+        
+        .form-group input:focus, .form-group textarea:focus {{
+            outline: none;
+            border-color: #007bff;
+            box-shadow: 0 0 0 2px rgba(0, 123, 255, 0.25);
+        }}
+        
+        .btn {{
+            background: #007bff;
+            color: white;
+            border: none;
+            padding: 10px 20px;
+            border-radius: 4px;
+            cursor: pointer;
+            font-size: 14px;
+            margin-right: 10px;
+            margin-bottom: 10px;
+        }}
+        
+        .btn:hover {{
+            background: #0056b3;
+        }}
+        
+        .btn:disabled {{
+            background: #6c757d;
+            cursor: not-allowed;
+        }}
+        
+        .btn-secondary {{
+            background: #6c757d;
+        }}
+        
+        .btn-secondary:hover {{
+            background: #545b62;
+        }}
+        
+        .state-display {{
+            background: white;
+            border: 1px solid #e0e0e0;
+            border-radius: 8px;
+            padding: 15px;
+            margin-bottom: 20px;
+        }}
+        
+        .state-item {{
+            margin-bottom: 8px;
+        }}
+        
+        .state-label {{
+            font-weight: 500;
+            color: #666;
+        }}
+        
+        .state-value {{
+            color: #333;
+            font-family: monospace;
+        }}
+        
+        .logs-container {{
+            background: white;
+            border: 1px solid #e0e0e0;
+            border-radius: 8px;
+            padding: 15px;
+            max-height: 400px;
+            overflow-y: auto;
+        }}
+        
+        .log-entry {{
+            border-bottom: 1px solid #f0f0f0;
+            padding: 10px 0;
+        }}
+        
+        .log-entry:last-child {{
+            border-bottom: none;
+        }}
+        
+        .log-timestamp {{
+            font-size: 12px;
+            color: #666;
+            margin-bottom: 5px;
+        }}
+        
+        .log-action {{
+            background: #e3f2fd;
+            padding: 8px;
+            border-radius: 4px;
+            margin-bottom: 5px;
+            font-family: monospace;
+            font-size: 12px;
+        }}
+        
+        .log-observation {{
+            background: #f3e5f5;
+            padding: 8px;
+            border-radius: 4px;
+            font-family: monospace;
+            font-size: 12px;
+        }}
+        
+        .log-reward {{
+            font-weight: 600;
+            color: #28a745;
+        }}
+        
+        .log-done {{
+            font-weight: 600;
+            color: #dc3545;
+        }}
+        
+        .status-indicator {{
+            display: inline-block;
+            width: 8px;
+            height: 8px;
+            border-radius: 50%;
+            margin-right: 8px;
+        }}
+        
+        .status-connected {{
+            background: #28a745;
+        }}
+        
+        .status-disconnected {{
+            background: #dc3545;
+        }}
+        
+        .json-display {{
+            background: #f8f9fa;
+            border: 1px solid #e9ecef;
+            border-radius: 4px;
+            padding: 10px;
+            font-family: monospace;
+            font-size: 12px;
+            white-space: pre-wrap;
+            max-height: 200px;
+            overflow-y: auto;
+        }}
+    </style>
+</head>
+<body>
+    <div class="container">
+        <!-- Left Pane: HumanAgent Interface -->
+        <div class="left-pane">
+            <div class="pane-header">
+                <span class="status-indicator status-disconnected" id="connection-status"></span>
+                HumanAgent Interface
+            </div>
+            <div class="pane-content">
+                <!-- Action Form -->
+                <div class="action-form">
+                    <h3>Take Action</h3>
+                    <form id="action-form">
+                        {_generate_action_form_fields(action_fields)}
+                        <button type="submit" class="btn" id="step-btn">Step</button>
+                    </form>
+                </div>
+                
+                <!-- Control Buttons -->
+                <div style="margin-bottom: 20px;">
+                    <button class="btn btn-secondary" id="reset-btn">Reset Environment</button>
+                    <button class="btn btn-secondary" id="state-btn">Get State</button>
+                </div>
+                
+                <!-- Current State Display -->
+                <div class="state-display">
+                    <h3>Current State</h3>
+                    <div id="current-state">
+                        <div class="state-item">
+                            <span class="state-label">Status:</span>
+                            <span class="state-value" id="env-status">Not initialized</span>
+                        </div>
+                        <div class="state-item">
+                            <span class="state-label">Episode ID:</span>
+                            <span class="state-value" id="episode-id">-</span>
+                        </div>
+                        <div class="state-item">
+                            <span class="state-label">Step Count:</span>
+                            <span class="state-value" id="step-count">0</span>
+                        </div>
+                    </div>
+                </div>
+            </div>
+        </div>
+        
+        <!-- Right Pane: State Observer -->
+        <div class="right-pane">
+            <div class="pane-header">
+                State Observer
+            </div>
+            <div class="pane-content">
+                <!-- Current Observation -->
+                <div class="state-display">
+                    <h3>Current Observation</h3>
+                    <div id="current-observation" class="json-display">
+                        No observation yet
+                    </div>
+                </div>
+                
+                <!-- Action Logs -->
+                <div class="logs-container">
+                    <h3>Action History</h3>
+                    <div id="action-logs">
+                        No actions taken yet
+                    </div>
+                </div>
+            </div>
+        </div>
+    </div>
+
+    <script>
+        class OpenEnvWebInterface {{
+            constructor() {{
+                this.ws = null;
+                this.isConnected = false;
+                this.init();
+            }}
+            
+            init() {{
+                this.connectWebSocket();
+                this.setupEventListeners();
+            }}
+            
+            connectWebSocket() {{
+                const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
+                const wsUrl = `${{protocol}}//${{window.location.host}}/ws`;
+                
+                this.ws = new WebSocket(wsUrl);
+                
+                this.ws.onopen = () => {{
+                    this.isConnected = true;
+                    this.updateConnectionStatus(true);
+                    console.log('WebSocket connected');
+                }};
+                
+                this.ws.onmessage = (event) => {{
+                    const data = JSON.parse(event.data);
+                    if (data.type === 'state_update') {{
+                        this.updateUI(data.episode_state);
+                    }}
+                }};
+                
+                this.ws.onclose = () => {{
+                    this.isConnected = false;
+                    this.updateConnectionStatus(false);
+                    console.log('WebSocket disconnected');
+                    // Attempt to reconnect after 3 seconds
+                    setTimeout(() => this.connectWebSocket(), 3000);
+                }};
+                
+                this.ws.onerror = (error) => {{
+                    console.error('WebSocket error:', error);
+                }};
+            }}
+            
+            setupEventListeners() {{
+                // Action form submission
+                document.getElementById('action-form').addEventListener('submit', (e) => {{
+                    e.preventDefault();
+                    this.submitAction();
+                }});
+                
+                // Reset button
+                document.getElementById('reset-btn').addEventListener('click', () => {{
+                    this.resetEnvironment();
+                }});
+                
+                // State button
+                document.getElementById('state-btn').addEventListener('click', () => {{
+                    this.getState();
+                }});
+            }}
+            
+            async submitAction() {{
+                const formData = new FormData(document.getElementById('action-form'));
+                const action = {{}};
+                
+                // Collect form data
+                for (const [key, value] of formData.entries()) {{
+                    if (value !== '') {{
+                        action[key] = value;
+                    }}
+                }}
+                
+                try {{
+                    const response = await fetch('/web/step', {{
+                        method: 'POST',
+                        headers: {{ 'Content-Type': 'application/json' }},
+                        body: JSON.stringify({{ action }})
+                    }});
+                    
+                    if (!response.ok) {{
+                        throw new Error(`HTTP error! status: ${{response.status}}`);
+                    }}
+                    
+                    const result = await response.json();
+                    console.log('Step result:', result);
+                }} catch (error) {{
+                    console.error('Error submitting action:', error);
+                    alert('Error submitting action: ' + error.message);
+                }}
+            }}
+            
+            async resetEnvironment() {{
+                try {{
+                    const response = await fetch('/web/reset', {{
+                        method: 'POST',
+                        headers: {{ 'Content-Type': 'application/json' }}
+                    }});
+                    
+                    if (!response.ok) {{
+                        throw new Error(`HTTP error! status: ${{response.status}}`);
+                    }}
+                    
+                    const result = await response.json();
+                    console.log('Reset result:', result);
+                }} catch (error) {{
+                    console.error('Error resetting environment:', error);
+                    alert('Error resetting environment: ' + error.message);
+                }}
+            }}
+            
+            async getState() {{
+                try {{
+                    const response = await fetch('/web/state');
+                    const state = await response.json();
+                    console.log('Current state:', state);
+                    alert('Current state: ' + JSON.stringify(state, null, 2));
+                }} catch (error) {{
+                    console.error('Error getting state:', error);
+                    alert('Error getting state: ' + error.message);
+                }}
+            }}
+            
+            updateConnectionStatus(connected) {{
+                const indicator = document.getElementById('connection-status');
+                if (connected) {{
+                    indicator.className = 'status-indicator status-connected';
+                }} else {{
+                    indicator.className = 'status-indicator status-disconnected';
+                }}
+            }}
+            
+            updateUI(episodeState) {{
+                // Update current state
+                document.getElementById('env-status').textContent = 
+                    episodeState.is_reset ? 'Reset' : 'Running';
+                document.getElementById('episode-id').textContent = 
+                    episodeState.episode_id || '-';
+                document.getElementById('step-count').textContent = 
+                    episodeState.step_count.toString();
+                
+                // Update current observation
+                const observationDiv = document.getElementById('current-observation');
+                if (episodeState.current_observation) {{
+                    observationDiv.textContent = JSON.stringify(
+                        episodeState.current_observation, null, 2
+                    );
+                }} else {{
+                    observationDiv.textContent = 'No observation yet';
+                }}
+                
+                // Update action logs
+                const logsDiv = document.getElementById('action-logs');
+                if (episodeState.action_logs.length === 0) {{
+                    logsDiv.innerHTML = 'No actions taken yet';
+                }} else {{
+                    logsDiv.innerHTML = episodeState.action_logs.map(log => `
+                        <div class="log-entry">
+                            <div class="log-timestamp">${{log.timestamp}} (Step ${{log.step_count}})</div>
+                            <div class="log-action">Action: ${{JSON.stringify(log.action, null, 2)}}</div>
+                            <div class="log-observation">Observation: ${{JSON.stringify(log.observation, null, 2)}}</div>
+                            <div>
+                                <span class="log-reward">Reward: ${{log.reward !== null ? log.reward : 'None'}}</span>
+                                ${{log.done ? '<span class="log-done">DONE</span>' : ''}}
+                            </div>
+                        </div>
+                    `).join('');
+                }}
+            }}
+        }}
+        
+        // Initialize the web interface when the page loads
+        document.addEventListener('DOMContentLoaded', () => {{
+            new OpenEnvWebInterface();
+        }});
+    </script>
+</body>
+</html>
+    """.replace('{_generate_action_form_fields(action_fields)}', _generate_action_form_fields(action_fields))
+
+
+def _generate_action_form_fields(action_fields: List[Dict[str, Any]]) -> str:
+    """Generate HTML form fields for action input."""
+    if not action_fields:
+        return '<p>No action fields available</p>'
+    
+    fields_html = []
+    for field in action_fields:
+        if field['type'] == 'checkbox':
+            fields_html.append(f'''
+                <div class="form-group">
+                    <label>
+                        <input type="checkbox" name="{field['name']}" value="true">
+                        {field['name']}
+                    </label>
+                </div>
+            ''')
+        elif field['type'] == 'text' and 'message' in field['name'].lower():
+            fields_html.append(f'''
+                <div class="form-group">
+                    <label for="{field['name']}">{field['name']}:</label>
+                    <textarea name="{field['name']}" id="{field['name']}" rows="3" placeholder="Enter {field['name']}..."></textarea>
+                </div>
+            ''')
+        else:
+            fields_html.append(f'''
+                <div class="form-group">
+                    <label for="{field['name']}">{field['name']}:</label>
+                    <input type="{field['type']}" name="{field['name']}" id="{field['name']}" placeholder="Enter {field['name']}..." {"required" if field['required'] else ""}>
+                </div>
+            ''')
+    
+    return '\n'.join(fields_html)

src/core/http_env_client.py

@@ -0,0 +1,175 @@
+"""
+core/runner_env.py
+Minimal HTTP-based environment client.
+- Talks to a single env worker exposing: POST /reset, POST /step
+
+Future hooks (commented below) for:
+- episode_id, seed on reset
+- request_id on step
+- custom headers (auth/trace)
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Dict, Generic, Optional, Type, TypeVar
+from .containers.runtime import LocalDockerProvider
+import requests
+
+from .types import StepResult
+
+if TYPE_CHECKING:
+    from .containers.runtime import ContainerProvider
+
+ActT = TypeVar("ActT")
+ObsT = TypeVar("ObsT")
+EnvClientT = TypeVar("EnvClientT", bound="HTTPEnvClient")
+
+
+class HTTPEnvClient(ABC, Generic[ActT, ObsT]):
+    def __init__(
+        self,
+        base_url: str,
+        request_timeout_s: float = 15.0,
+        default_headers: Optional[Dict[str, str]] = None,
+        provider: Optional["ContainerProvider"] = None,
+    ):
+        self._base = base_url.rstrip("/")
+        self._timeout = float(request_timeout_s)
+        self._http = requests.Session()
+        self._headers = default_headers or {}
+        self._provider = provider
+
+    @classmethod
+    def from_docker_image(
+        cls: Type[EnvClientT],
+        image: str,
+        provider: Optional["ContainerProvider"] = None,
+    ) -> EnvClientT:
+        """
+        Create an environment client by spinning up a Docker container locally.
+
+        This is a development utility that:
+        1. Starts a Docker container from the specified image
+        2. Waits for the server to be ready
+        3. Creates and returns a client instance connected to the container
+
+        Note: The container lifecycle management is left to the user or higher-level
+        orchestration. The container will keep running until manually stopped.
+
+        Args:
+            image: Docker image name to run (e.g., "echo-env:latest")
+            provider: Container provider to use (defaults to LocalDockerProvider)
+
+        Returns:
+            An instance of the client class connected to the running container
+
+        Example:
+            >>> from envs.coding_env.client import CodingEnv
+            >>> from envs.coding_env.models import CodeAction
+            >>>
+            >>> # Create environment from image
+            >>> env = CodingEnv.from_docker_image("coding-env:latest")
+            >>>
+            >>> # Use the environment
+            >>> result = env.reset()
+            >>> print(result.observation)
+            >>>
+            >>> step_result = env.step(CodeAction(code="print('hello')"))
+            >>> print(step_result.observation.stdout)
+            >>>
+            >>> # Cleanup (optional)
+            >>> env.close()
+        """
+
+        # Use default provider if none provided
+        if provider is None:
+            provider = LocalDockerProvider()
+
+        # 1. Start container
+        base_url = provider.start_container(image)
+
+        # 2. Wait for server to be ready
+        provider.wait_for_ready(base_url)
+
+        # 3. Create and return client instance with provider reference
+        return cls(base_url=base_url, provider=provider)
+
+    @abstractmethod
+    def _step_payload(self, action: ActT) -> dict:
+        """Convert an Action object to the JSON body expected by the env server."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _parse_result(self, payload: dict) -> StepResult[ObsT]:
+        """Convert a JSON response from the env server to StepResult[ObsT]."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _parse_state(self, payload: dict) -> Any:
+        """Convert a JSON response from the state endpoint to a State object."""
+        raise NotImplementedError
+
+    # ---------- Environment Server Interface Methods ----------
+    def reset(self) -> StepResult[ObsT]:
+        body: Dict[str, Any] = {}
+        # TODO: later:
+        # body["seed"] = seed
+        # body["episode_id"] = episode_id
+        r = self._http.post(
+            f"{self._base}/reset",
+            json=body,
+            headers=self._headers,
+            timeout=self._timeout,
+        )
+        r.raise_for_status()
+        return self._parse_result(r.json())
+
+    def step(self, action: ActT) -> StepResult[ObsT]:
+        body: Dict[str, Any] = {
+            "action": self._step_payload(action),
+            "timeout_s": int(self._timeout),
+        }
+        # TODO: later:
+        # body["request_id"] = str(uuid.uuid4())
+        # body["episode_id"] = current_episode_id
+        r = self._http.post(
+            f"{self._base}/step",
+            json=body,
+            headers=self._headers,
+            timeout=self._timeout,
+        )
+        r.raise_for_status()
+        return self._parse_result(r.json())
+
+    def state(self) -> Any:
+        """
+        Get the current environment state from the server.
+
+        Returns:
+            State object with environment state information (e.g., episode_id, step_count)
+
+        Example:
+            >>> client = EchoEnv.from_docker_image("echo-env:latest")
+            >>> result = client.reset()
+            >>> state = client.state()
+            >>> print(state.episode_id)
+            >>> print(state.step_count)
+        """
+        r = self._http.get(
+            f"{self._base}/state",
+            headers=self._headers,
+            timeout=self._timeout,
+        )
+        r.raise_for_status()
+        return self._parse_state(r.json())
+
+    def close(self) -> None:
+        """
+        Close the environment and clean up resources.
+
+        If this client was created via from_docker_image(), this will stop
+        and remove the associated container.
+        """
+        if self._provider is not None:
+            self._provider.stop_container()

src/core/tools/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Core tools for code execution and other utilities."""
+
+from .local_python_executor import PyExecutor
+
+__all__ = ["PyExecutor"]

src/core/tools/local_python_executor.py

@@ -0,0 +1,105 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Local Python Executor.
+
+This module provides functionality for executing Python code locally by wrapping
+the smolagents LocalPythonExecutor.
+"""
+
+from smolagents import LocalPythonExecutor
+
+from core.env_server.types import CodeExecResult
+
+
+class PyExecutor:
+    """
+    Wrapper around smolagents LocalPythonExecutor for executing Python code.
+
+    This class provides a simple interface to execute Python code in a subprocess
+    and capture the results including stdout, stderr, and exit code.
+
+    Args:
+        additional_imports: List of additional module imports to authorize.
+                          For example: ["numpy", "pandas", "matplotlib"]
+                          These will be added to the base authorized imports.
+
+    Example:
+        >>> # Basic usage with default imports
+        >>> executor = PyExecutor()
+        >>> result = executor.run("print('Hello, World!')")
+        >>> print(result.stdout)  # "Hello, World!\n"
+        >>> print(result.exit_code)  # 0
+        >>>
+        >>> # Usage with additional imports
+        >>> executor = PyExecutor(additional_imports=["numpy", "pandas"])
+        >>> result = executor.run("import numpy as np\\nprint(np.array([1, 2, 3]))")
+        >>> print(result.stdout)  # "[1 2 3]\n"
+    """
+
+    def __init__(self, additional_imports: list[str] | None = None):
+        """
+        Initialize the PyExecutor with a LocalPythonExecutor instance.
+
+        Args:
+            additional_imports: List of additional module names to authorize for import.
+                              Defaults to an empty list if not provided.
+        """
+        if additional_imports is None:
+            additional_imports = []
+        self._executor = LocalPythonExecutor(
+            additional_authorized_imports=additional_imports
+        )
+        # Initialize tools to make BASE_PYTHON_TOOLS available (including print)
+        self._executor.send_tools({})
+
+    def run(self, code: str) -> CodeExecResult:
+        """
+        Execute Python code and return the result.
+
+        Args:
+            code: Python code string to execute
+
+        Returns:
+            CodeExecResult containing stdout, stderr, and exit_code
+
+        Example:
+            >>> executor = PyExecutor()
+            >>> result = executor.run("x = 5 + 3\\nprint(x)")
+            >>> print(result.stdout)  # "8\n"
+            >>> print(result.exit_code)  # 0
+            >>>
+            >>> # Error handling
+            >>> result = executor.run("1 / 0")
+            >>> print(result.exit_code)  # 1
+            >>> print(result.stderr)  # Contains error message
+        """
+        try:
+            # Execute the code using LocalPythonExecutor
+            # LocalPythonExecutor returns a CodeOutput object with output, logs, is_final_answer
+            exec_result = self._executor(code)
+
+            # Extract the logs (which contain print outputs) as stdout
+            # The output field contains the return value of the code
+            stdout = exec_result.logs
+            stderr = ""
+            exit_code = 0  # Success
+
+            return CodeExecResult(
+                stdout=stdout,
+                stderr=stderr,
+                exit_code=exit_code,
+            )
+
+        except Exception as e:
+            # LocalPythonExecutor raises InterpreterError for various issues
+            # (syntax errors, forbidden operations, runtime errors, etc.)
+            return CodeExecResult(
+                stdout="",
+                stderr=str(e),
+                exit_code=1,  # Non-zero indicates error
+            )

src/core/types.py

@@ -0,0 +1,22 @@
+# Type definitions for EnvTorch
+from dataclasses import dataclass
+from typing import Any, Generic, Optional, TypeVar
+
+# Generic type for observations
+ObsT = TypeVar("ObsT")  # TypeVar for typehinting in IDEs
+
+
+@dataclass
+class StepResult(Generic[ObsT]):
+    """
+    Represents the result of one environment step.
+
+    Attributes:
+        observation: The environment's observation after the action.
+        reward: Scalar reward for this step (optional).
+        done: Whether the episode is finished.
+    """
+
+    observation: ObsT
+    reward: Optional[float] = None
+    done: bool = False

src/envs/chat_env/README.md

@@ -0,0 +1,268 @@
+# Chat Environment
+
+A 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.
+
+## Overview
+
+ChatEnvironment is a lightweight environment that:
+- Manages conversation history in Huggingface chat format
+- Handles tokenization internally using any compatible tokenizer
+- Stores both messages and tokens for efficient model interaction
+- Provides a clean interface for building chat-based RL agents
+
+ChatEnvironment can be used in **two ways**:
+1. **Direct usage**: Import and use ChatEnvironment directly in your Python code (best for local development)
+2. **HTTP client**: Use ChatEnv client to connect to a ChatEnvironment server (best for distributed/containerized deployments)
+
+## Quick Start
+
+### Option 1: Direct Usage (Local)
+
+```python
+from transformers import AutoTokenizer
+from envs.chat_env import ChatAction, ChatObservation
+from envs.chat_env.server import ChatEnvironment
+from core.env_server import Message
+
+# Initialize with a tokenizer and optional system prompt
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+env = ChatEnvironment(
+    tokenizer=tokenizer,
+    system_prompt="You are a helpful assistant.",
+    system_role="system"
+)
+
+# Reset the environment
+obs = env.reset()
+print(f"Messages: {obs.messages}")
+print(f"Tokens shape: {obs.tokens.shape}")
+
+# Create an action from a message
+user_message: Message = {"role": "user", "content": "Hello!"}
+action = env.message_to_action(user_message)
+
+# Step the environment
+obs = env.step(action)
+print(f"Updated messages: {obs.messages}")
+print(f"Updated tokens shape: {obs.tokens.shape}")
+```
+
+### Option 2: HTTP Client (Distributed)
+
+```python
+from transformers import AutoTokenizer
+from envs.chat_env import ChatEnv, ChatAction
+import torch
+
+# Create environment from Docker image
+client = ChatEnv.from_docker_image("chat-env:latest")
+
+# Or connect to existing server
+# client = ChatEnv(base_url="http://localhost:8000")
+
+# Reset
+result = client.reset()
+print(f"Initial messages: {result.observation.messages}")
+
+# Send an action with tokens
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+message = {"role": "user", "content": "Hello!"}
+action = client.message_to_action(message, tokenizer)
+
+result = client.step(action)
+print(f"Messages: {result.observation.messages}")
+print(f"Reward: {result.reward}")
+
+# Cleanup
+client.close()
+```
+
+### Building the Docker Image
+
+Before using the HTTP client, build the Docker image:
+
+```bash
+# From project root
+docker build -t chat-env:latest -f src/envs/chat_env/server/Dockerfile .
+
+# Optionally specify a different tokenizer
+docker build -t chat-env:latest \
+  --build-arg TOKENIZER_NAME=meta-llama/Llama-2-7b-chat-hf \
+  -f src/envs/chat_env/server/Dockerfile .
+```
+
+## Architecture
+
+### Data Models
+
+#### ChatAction
+Actions contain only tokens (PyTorch tensors) that interface directly with models:
+```python
+@dataclass
+class ChatAction(Action):
+    tokens: torch.Tensor  # Required, cannot be empty
+```
+
+#### ChatObservation
+Observations contain both the message history and flattened tokens:
+```python
+@dataclass
+class ChatObservation(Observation):
+    messages: list[Message]  # List of {"role": str, "content": str}
+    tokens: torch.Tensor     # Flattened tensor of all conversation tokens
+    # Inherited: done, reward, metadata
+```
+
+#### ChatState
+Internal state tracking message and token history:
+```python
+@dataclass
+class ChatState(State):
+    history_messages: list[Message]
+    history_tokens: list[torch.Tensor]
+    # Inherited: episode_id, step_count
+```
+
+### Key Methods
+
+#### `reset() -> ChatObservation`
+Resets the environment to initial state with optional system prompt.
+
+#### `step(action: ChatAction) -> ChatObservation`
+Takes an action (tokens), decodes to text, adds to history, returns updated observation.
+
+#### `message_to_action(message: Message) -> ChatAction`
+Convenience method to convert a message dict to a tokenized ChatAction.
+
+## Usage Patterns
+
+### Basic Conversation
+
+```python
+from transformers import AutoTokenizer
+from envs.chat_env.server import ChatEnvironment
+from core.env_server import Message
+
+tokenizer = AutoTokenizer.from_pretrained("gpt2")
+env = ChatEnvironment(tokenizer=tokenizer)
+
+# Reset
+obs = env.reset()
+
+# User turn
+user_msg: Message = {"role": "user", "content": "What is 2+2?"}
+action = env.message_to_action(user_msg)
+obs = env.step(action)
+
+# Assistant turn
+assistant_msg: Message = {"role": "assistant", "content": "2+2 equals 4."}
+action = env.message_to_action(assistant_msg)
+obs = env.step(action)
+
+# Access conversation history
+print(f"Full conversation: {obs.messages}")
+print(f"All tokens: {obs.tokens}")
+```
+
+### With Transforms
+
+You can add transforms to compute rewards or modify observations:
+
+```python
+from core.env_server import Transform, Observation
+
+class LengthRewardTransform(Transform):
+    """Reward based on response length."""
+
+    def __call__(self, observation: Observation) -> Observation:
+        if hasattr(observation, 'messages') and observation.messages:
+            last_message = observation.messages[-1]
+            observation.reward = len(last_message['content']) * 0.1
+        return observation
+
+env = ChatEnvironment(
+    tokenizer=tokenizer,
+    transform=LengthRewardTransform()
+)
+```
+
+### Direct Token Usage
+
+If you're generating tokens from a model, you can create actions directly:
+
+```python
+import torch
+from envs.chat_env import ChatAction
+
+# Assume you have tokens from your model
+generated_tokens = torch.tensor([[1, 2, 3, 4, 5]])
+
+# Create action directly
+action = ChatAction(tokens=generated_tokens)
+
+# Step environment
+obs = env.step(action)
+```
+
+## Design Philosophy
+
+ChatEnvironment is intentionally minimal and flexible:
+
+1. **No HTTP overhead**: Works directly with Python objects and tensors
+2. **Tokenizer ownership**: Environment handles tokenization consistently
+3. **Dual representation**: Maintains both human-readable messages and model-ready tokens
+4. **Transform support**: Extensible reward computation and observation modification
+5. **Type-safe**: Uses typed Messages compatible with Huggingface format
+
+## Integration with Models
+
+ChatEnvironment pairs naturally with language models:
+
+```python
+# Pseudo-code for RL training loop
+model = YourLanguageModel()
+env = ChatEnvironment(tokenizer=model.tokenizer)
+
+for episode in range(num_episodes):
+    obs = env.reset()
+
+    while not obs.done:
+        # Model generates response tokens
+        action_tokens = model.generate(obs.tokens)
+        action = ChatAction(tokens=action_tokens)
+
+        # Step environment
+        obs = env.step(action)
+
+        # Use obs.reward for RL updates
+        model.update(obs.reward)
+```
+
+## Project Structure
+
+```
+chat_env/
+├── __init__.py              # Module exports (ChatEnv, ChatAction, etc.)
+├── README.md                # This file
+├── client.py                # ChatEnv HTTP client
+├── models.py                # ChatAction, ChatObservation, ChatState
+└── server/
+    ├── __init__.py          # Server module exports
+    ├── chat_environment.py  # Core ChatEnvironment implementation
+    ├── app.py               # FastAPI server application
+    ├── test_chat_env.py     # Unit tests
+    └── Dockerfile           # Container image for HTTP server
+```
+
+## Requirements
+
+- Python 3.10+
+- PyTorch
+- A tokenizer with `apply_chat_template` method (e.g., Huggingface transformers)
+
+## Notes
+
+- ChatEnvironment does **not** generate responses - it only manages conversation state
+- You need to provide tokens from your model or other source
+- The environment is thread-safe for single-threaded use only
+- For multi-turn conversations, alternate between user and assistant messages

src/envs/chat_env/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Chat Environment - A chat-based environment for LLMs with tokenization support."""
+
+from .client import ChatEnv
+from .models import ChatAction, ChatObservation, ChatState
+
+__all__ = ["ChatAction", "ChatObservation", "ChatState", "ChatEnv"]

src/envs/chat_env/client.py

@@ -0,0 +1,182 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Chat Environment HTTP Client.
+
+This module provides the client for connecting to a Chat Environment server
+over HTTP.
+"""
+
+from typing import Any, Dict
+
+import torch
+
+from core.env_server.interfaces import Message
+from core.env_server.types import State
+from core.http_env_client import HTTPEnvClient
+from core.types import StepResult
+
+from .models import ChatAction, ChatObservation, ChatState
+
+
+class ChatEnv(HTTPEnvClient[ChatAction, ChatObservation]):
+    """
+    HTTP client for the Chat Environment.
+
+    This client connects to a ChatEnvironment HTTP server and provides
+    methods to interact with it: reset(), step(), and state access.
+
+    Note: Since ChatEnvironment works with PyTorch tensors, the HTTP layer
+    serializes tokens as lists for transport and deserializes them back to tensors.
+
+    Example:
+        >>> # Connect to a running server
+        >>> client = ChatEnv(base_url="http://localhost:8000")
+        >>> result = client.reset()
+        >>> print(result.observation.messages)
+        >>>
+        >>> # Send an action with tokens
+        >>> import torch
+        >>> tokens = torch.tensor([[1, 2, 3, 4, 5]])
+        >>> result = client.step(ChatAction(tokens=tokens))
+        >>> print(result.observation.messages)
+        >>> print(result.reward)
+
+    Example with Docker:
+        >>> # Automatically start container and connect
+        >>> client = ChatEnv.from_docker_image("chat-env:latest")
+        >>> result = client.reset()
+        >>> result = client.step(ChatAction(tokens=torch.tensor([[1, 2, 3]])))
+    """
+
+    def _step_payload(self, action: ChatAction) -> Dict:
+        """
+        Convert ChatAction to JSON payload for step request.
+
+        Since PyTorch tensors can't be directly serialized to JSON,
+        we convert them to nested lists.
+
+        Args:
+            action: ChatAction instance with tokens
+
+        Returns:
+            Dictionary representation suitable for JSON encoding
+        """
+        # Convert tensor to list for JSON serialization
+        if isinstance(action.tokens, torch.Tensor):
+            tokens_list = action.tokens.tolist()
+        else:
+            tokens_list = action.tokens
+
+        return {
+            "tokens": tokens_list,
+            "metadata": action.metadata,
+        }
+
+    def _parse_result(self, payload: Dict) -> StepResult[ChatObservation]:
+        """
+        Parse server response into StepResult[ChatObservation].
+
+        Args:
+            payload: JSON response from server
+
+        Returns:
+            StepResult with ChatObservation
+        """
+        obs_data = payload.get("observation", {})
+
+        # Convert tokens list back to tensor
+        tokens_data = obs_data.get("tokens", [])
+        if isinstance(tokens_data, list):
+            if tokens_data:
+                tokens = torch.tensor(tokens_data)
+            else:
+                tokens = torch.tensor([])
+        else:
+            tokens = torch.tensor([])
+
+        # Parse messages
+        messages = obs_data.get("messages", [])
+
+        observation = ChatObservation(
+            messages=messages,
+            tokens=tokens,
+            done=payload.get("done", False),
+            reward=payload.get("reward"),
+            metadata=obs_data.get("metadata", {}),
+        )
+
+        return StepResult(
+            observation=observation,
+            reward=payload.get("reward"),
+            done=payload.get("done", False),
+        )
+
+    def _parse_state(self, payload: Dict) -> ChatState:
+        """
+        Parse server response into ChatState object.
+
+        Args:
+            payload: JSON response from /state endpoint
+
+        Returns:
+            ChatState object with conversation history
+        """
+        # Parse history messages
+        history_messages = payload.get("history_messages", [])
+
+        # Parse history tokens - convert lists back to tensors
+        history_tokens_data = payload.get("history_tokens", [])
+        history_tokens = []
+        for token_list in history_tokens_data:
+            if token_list:
+                history_tokens.append(torch.tensor(token_list))
+            else:
+                history_tokens.append(torch.tensor([]))
+
+        return ChatState(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+            history_messages=history_messages,
+            history_tokens=history_tokens,
+        )
+
+    def message_to_action(self, message: Message, tokenizer: Any) -> ChatAction:
+        """
+        Helper method to convert a message to a ChatAction using a tokenizer.
+
+        This is a client-side convenience method for users who have a tokenizer
+        and want to create actions from messages.
+
+        Args:
+            message: Message dict with 'role' and 'content'
+            tokenizer: Tokenizer with apply_chat_template method
+
+        Returns:
+            ChatAction with tokenized message
+
+        Example:
+            >>> from transformers import AutoTokenizer
+            >>> tokenizer = AutoTokenizer.from_pretrained("gpt2")
+            >>> client = ChatEnv(base_url="http://localhost:8000")
+            >>> message = {"role": "user", "content": "Hello!"}
+            >>> action = client.message_to_action(message, tokenizer)
+            >>> result = client.step(action)
+        """
+        if "role" not in message:
+            raise ValueError("Message must contain a 'role' key")
+        if "content" not in message:
+            raise ValueError("Message must contain a 'content' key")
+        if message["content"] is None:
+            raise ValueError("Message content cannot be None")
+
+        # Tokenize the message
+        tokens = tokenizer.apply_chat_template(
+            conversation=[message], tokenize=True, return_tensors="pt"
+        )
+
+        return ChatAction(tokens=tokens)

src/envs/chat_env/models.py

@@ -0,0 +1,67 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Data models for the Chat Environment.
+
+The Chat environment provides a chat-based interface for LLMs with support
+for tokenization and message history management.
+"""
+
+from dataclasses import dataclass, field
+
+import torch
+
+from core.env_server.interfaces import Message
+from core.env_server.types import Action, Observation, State
+
+
+@dataclass
+class ChatAction(Action):
+    """Action for chat environments.
+
+    Contains tokens that represent the action to be taken.
+    This interfaces directly with models.
+    """
+
+    tokens: torch.Tensor = field(default_factory=lambda: torch.tensor([]))
+
+    def __post_init__(self):
+        """Validate required fields after initialization."""
+        if self.tokens.numel() == 0:
+            raise ValueError("tokens is required and cannot be empty")
+
+
+@dataclass
+class ChatState(State):
+    """State of the ChatEnvironment containing message history."""
+
+    history_messages: list[Message] = field(default_factory=list)
+    history_tokens: list[torch.Tensor] = field(
+        default_factory=list
+    )  # Same len as messages
+
+
+@dataclass(kw_only=True)
+class ChatObservation(Observation):
+    """Observation returned by ChatEnvironment.
+
+    Contains the message history in Huggingface format (list of dicts with role/content)
+    and the tokenized representation of the entire conversation.
+
+    The environment owns the tokenizer and generates the tokens from the messages.
+
+    Example:
+    messages = [
+     {"role": "system", "content": "You are a helpful assistant"},
+     {"role": "user", "content": "How tall is the Eiffel Tower?"},
+    ]
+    tokens = tensor([1, 2, 3, 4, 5, ...])  # tokenized entire conversation
+    """
+
+    messages: list[Message] = field(default_factory=list)
+    tokens: torch.Tensor = field(default_factory=lambda: torch.tensor([]))
+    # Inherited fields from Observation ABC: reward, done, metadata

src/envs/chat_env/server/Dockerfile

@@ -0,0 +1,29 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# Use the standard openenv base image
+# Built from: docker build -t openenv-base:latest -f src/core/containers/images/Dockerfile .
+# In GitHub Actions, this is overridden to use the GHCR base image
+ARG BASE_IMAGE=openenv-base:latest
+FROM ${BASE_IMAGE}
+
+# Install additional dependencies for ChatEnvironment
+RUN pip install torch transformers
+
+# Copy only what's needed for this environment
+COPY src/core/ /app/src/core/
+COPY src/envs/chat_env/ /app/src/envs/chat_env/
+
+# Environment variables that can be overridden at runtime
+ENV TOKENIZER_NAME=gpt2
+ENV SYSTEM_PROMPT="You are a helpful AI assistant."
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=3s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:8000/health || exit 1
+
+# Run the FastAPI server
+CMD ["uvicorn", "envs.chat_env.server.app:app", "--host", "0.0.0.0", "--port", "8000"]

src/envs/chat_env/server/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Chat environment server components."""
+
+from .chat_environment import ChatEnvironment
+
+__all__ = ["ChatEnvironment"]

src/envs/chat_env/server/app.py

@@ -0,0 +1,77 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+FastAPI application for the Chat Environment.
+
+This module creates an HTTP server that exposes the ChatEnvironment
+over HTTP endpoints, making it compatible with HTTPEnvClient.
+
+Note: This server requires a tokenizer to be initialized. The tokenizer
+must be specified when starting the server.
+
+Usage:
+    # Development (with auto-reload):
+    uvicorn envs.chat_env.server.app:app --reload --host 0.0.0.0 --port 8000
+
+    # Production:
+    uvicorn envs.chat_env.server.app:app --host 0.0.0.0 --port 8000 --workers 4
+
+    # Or run directly:
+    python -m envs.chat_env.server.app
+"""
+
+import os
+
+from core.env_server import create_app
+
+from ..models import ChatAction, ChatObservation
+from .chat_environment import ChatEnvironment
+
+
+# Initialize tokenizer based on environment variable
+def get_tokenizer():
+    """Get tokenizer from environment or use a mock for testing."""
+    tokenizer_name = os.environ.get("TOKENIZER_NAME", "gpt2")
+
+    try:
+        from transformers import AutoTokenizer
+
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+        print(f"Loaded tokenizer: {tokenizer_name}")
+        return tokenizer
+    except ImportError:
+        print(
+            "Warning: transformers not installed, using mock tokenizer for testing only"
+        )
+        # Use mock tokenizer from tests
+        import sys
+        from pathlib import Path
+
+        # Add parent directory to path to import test utilities
+        test_path = Path(__file__).parent
+        sys.path.insert(0, str(test_path))
+
+        from test_chat_env import MockTokenizer
+
+        return MockTokenizer()
+
+
+# Get system prompt from environment
+system_prompt = os.environ.get("SYSTEM_PROMPT", None)
+
+# Create the environment instance with tokenizer
+tokenizer = get_tokenizer()
+env = ChatEnvironment(tokenizer=tokenizer, system_prompt=system_prompt)
+
+# Create the FastAPI app with routes
+app = create_app(env, ChatAction, ChatObservation)
+
+
+if __name__ == "__main__":
+    import uvicorn
+
+    uvicorn.run(app, host="0.0.0.0", port=8000)

src/envs/chat_env/server/chat_environment.py

@@ -0,0 +1,172 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Chat Environment Implementation.
+
+A chat-based environment for LLMs, designed as a blank canvas for conversation and RL.
+"""
+
+import torch
+
+from core.env_server.interfaces import Environment, Message, ModelTokenizer, Transform
+
+from ..models import ChatAction, ChatObservation, ChatState
+
+
+class ChatEnvironment(Environment):
+    """A chat-based environment for LLMs, designed as a blank canvas for conversation and RL.
+
+    This environment is designed to work with language models. It provides the fundamental structure
+    for managing conversation state but is intentionally minimal to allow maximum flexibility.
+
+    The environment owns the tokenizer and is responsible for managing both message history and tokens.
+    Actions contain only tokens that interface directly with models.
+
+    Args:
+        tokenizer: A tokenizer that will be used to tokenize the conversation
+        system_prompt: An optional system prompt string to use during reset calls (optional)
+        system_role: The role of the system (at reset time). Defaults to "system"
+        transform: Optional transform to apply to observations
+    """
+
+    def __init__(
+        self,
+        tokenizer: ModelTokenizer,
+        system_prompt: str | None = None,
+        system_role: str = "system",
+        transform: Transform | None = None,
+    ):
+        super().__init__(transform=transform)
+
+        if not hasattr(tokenizer, "apply_chat_template"):
+            raise ValueError("Tokenizer must have 'apply_chat_template' method")
+        self.tokenizer = tokenizer
+        self.system_prompt = system_prompt
+        self.system_role = system_role
+
+        self._state = ChatState()
+
+        if system_prompt:
+            system_message: Message = {"role": system_role, "content": system_prompt}
+            self._state.history_messages.append(system_message)
+            # Tokenize the system message
+            system_tokens = self.tokenizer.apply_chat_template(
+                conversation=[system_message], tokenize=True, return_tensors="pt"  # type: ignore
+            )
+            self._state.history_tokens.append(system_tokens)
+
+    def reset(self) -> ChatObservation:
+        """Reset the environment to initial state.
+
+        Returns:
+            ChatObservation: Initial observation with system prompt (if any)
+        """
+        self._state.history_messages = []
+        self._state.history_tokens = []
+        if self.system_prompt:
+            system_message: Message = {
+                "role": self.system_role,
+                "content": self.system_prompt,
+            }
+            self._state.history_messages = [system_message]
+            # Tokenize the system message
+            system_tokens = self.tokenizer.apply_chat_template(
+                conversation=[system_message], tokenize=True, return_tensors="pt"  # type: ignore
+            )
+            self._state.history_tokens = [system_tokens]
+
+        return self._create_observation()
+
+    def step(self, action: ChatAction) -> ChatObservation:  # type: ignore[override]
+        """Take a step in the environment by adding tokens to the chat history.
+
+        Args:
+            action: A ChatAction object containing tokens.
+
+        Returns:
+            ChatObservation: The updated observation with the new tokens added.
+        """
+        # Store the tokens directly from the action
+        self._state.history_tokens.append(action.tokens)
+
+        # Decode tokens to text and add as a message to history
+        decoded_text = self.tokenizer.decode(
+            action.tokens.squeeze(), skip_special_tokens=True
+        )
+        assistant_message: Message = {"role": "assistant", "content": decoded_text}
+        self._state.history_messages.append(assistant_message)
+
+        return self._create_observation()
+
+    def _create_observation(self) -> ChatObservation:
+        """Create a ChatObservation from the current state.
+
+        Returns both the message history and the tokens flattened as a single tensor
+        ready to be used by models.
+
+        Returns:
+            ChatObservation: Observation with messages and flattened tokens
+        """
+        if self._state.history_tokens:
+            # Flatten all tokens into a single 1D tensor
+            flattened_tokens = torch.cat(
+                (t.flatten() for t in self._state.history_tokens), dim=0
+            )
+        else:
+            flattened_tokens = torch.tensor([])
+
+        observation = ChatObservation(
+            messages=self._state.history_messages.copy(),  # Copy to prevent external mutation
+            tokens=flattened_tokens,
+        )
+
+        transformed = self._apply_transform(observation)
+        if isinstance(transformed, ChatObservation):
+            return transformed
+        else:
+            # If transform returns base Observation, convert back to ChatObservation
+            return ChatObservation(
+                messages=getattr(transformed, "messages", []),
+                tokens=getattr(transformed, "tokens", torch.tensor([])),
+                done=transformed.done,
+                reward=transformed.reward,
+            )
+
+    @property
+    def state(self) -> ChatState:
+        """Get the current state of the environment.
+
+        Returns:
+            ChatState: The current state.
+        """
+        return self._state
+
+    def message_to_action(self, message: Message) -> ChatAction:
+        """Convert a message dictionary to a ChatAction with tokens.
+
+        Args:
+            message: Dictionary with 'role' and 'content' keys
+
+        Returns:
+            ChatAction: A new ChatAction instance with tokenized content
+
+        Raises:
+            ValueError: If required keys are missing
+        """
+        if "role" not in message:
+            raise ValueError("Message must contain a 'role' key")
+        if "content" not in message:
+            raise ValueError("Message must contain a 'content' key")
+        if message["content"] is None:
+            raise ValueError("Message content cannot be None")
+
+        # Tokenize the single message
+        tokens = self.tokenizer.apply_chat_template(
+            conversation=[message], tokenize=True, return_tensors="pt"  # type: ignore
+        )
+
+        return ChatAction(tokens=tokens)

src/envs/chat_env/server/test_chat_env.py

@@ -0,0 +1,328 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Test suite for ChatEnvironment.
+
+Proper unit tests with assertions to verify correct behavior.
+"""
+
+import torch
+
+from core.env_server.interfaces import Message
+
+from ..models import ChatAction
+from .chat_environment import ChatEnvironment
+
+
+class MockTokenizer:
+    """Mock tokenizer for testing without requiring transformers library."""
+
+    def apply_chat_template(
+        self,
+        conversation: list[Message],
+        tokenize: bool = True,
+        return_tensors: str | None = None,
+        **kwargs,
+    ):
+        """Mock implementation that creates deterministic token tensors from text."""
+        # Concatenate all message content
+        text = " ".join([msg["content"] for msg in conversation])
+
+        # Create deterministic tokens based on text content
+        # Use character codes modulo 256 to get valid token IDs
+        tokens = [ord(c) % 256 for c in text]
+
+        if return_tensors == "pt":
+            return torch.tensor([tokens])
+        return tokens
+
+    def decode(self, token_ids, skip_special_tokens: bool = False, **kwargs) -> str:
+        """Mock decode that reverses the encoding process."""
+        if isinstance(token_ids, torch.Tensor):
+            token_ids = token_ids.tolist()
+
+        # Reverse the encoding: convert tokens back to characters
+        chars = [chr(t) for t in token_ids]
+        return "".join(chars)
+
+
+def test_tokenization_consistency():
+    """Test that tokenizing the same string produces the same tokens."""
+    tokenizer = MockTokenizer()
+    env = ChatEnvironment(tokenizer=tokenizer)
+
+    # Create the same message twice
+    message1: Message = {"role": "user", "content": "Hello, world!"}
+    message2: Message = {"role": "user", "content": "Hello, world!"}
+
+    # Convert to actions
+    action1 = env.message_to_action(message1)
+    action2 = env.message_to_action(message2)
+
+    # Verify tokens are identical
+    assert torch.equal(
+        action1.tokens, action2.tokens
+    ), "Same message should produce identical tokens"
+
+    # Verify tokens are not empty
+    assert action1.tokens.numel() > 0, "Tokens should not be empty"
+
+    print("✓ test_tokenization_consistency passed")
+
+
+def test_message_content_preservation():
+    """Test that message content is preserved in the observation."""
+    tokenizer = MockTokenizer()
+    env = ChatEnvironment(tokenizer=tokenizer)
+
+    env.reset()
+
+    # Test with user message
+    user_content = "What is the capital of France?"
+    user_message: Message = {"role": "user", "content": user_content}
+    action = env.message_to_action(user_message)
+    obs = env.step(action)
+
+    # The last message should have the decoded content
+    assert len(obs.messages) > 0, "Observation should have at least one message"
+    last_message = obs.messages[-1]
+
+    # Verify the decoded content matches what we sent
+    # Note: The environment decodes the tokens, so we verify the round-trip
+    decoded_content = last_message["content"]
+    assert decoded_content == user_content, (
+        f"Message content should be preserved. "
+        f"Expected: {user_content}, Got: {decoded_content}"
+    )
+
+    # Test with assistant message
+    assistant_content = "The capital of France is Paris."
+    assistant_message: Message = {"role": "assistant", "content": assistant_content}
+    action = env.message_to_action(assistant_message)
+    obs = env.step(action)
+
+    # Verify the last message has the assistant content
+    assert len(obs.messages) >= 2, "Should have at least 2 messages now"
+    last_message = obs.messages[-1]
+    decoded_content = last_message["content"]
+    assert decoded_content == assistant_content, (
+        f"Assistant message content should be preserved. "
+        f"Expected: {assistant_content}, Got: {decoded_content}"
+    )
+
+    print("✓ test_message_content_preservation passed")
+
+
+def test_system_prompt_preserved():
+    """Test that system prompt is preserved after reset."""
+    tokenizer = MockTokenizer()
+    system_prompt = "You are a helpful assistant."
+
+    env = ChatEnvironment(tokenizer=tokenizer, system_prompt=system_prompt)
+
+    # Check after initialization
+    obs = env.reset()
+    assert len(obs.messages) == 1, "Should have exactly one message (system prompt)"
+    assert obs.messages[0]["role"] == "system", "First message should have system role"
+    assert (
+        obs.messages[0]["content"] == system_prompt
+    ), "System prompt content should match"
+
+    # Add some messages
+    action = env.message_to_action({"role": "user", "content": "Hello"})
+    env.step(action)
+
+    # Reset and verify system prompt is still there
+    obs = env.reset()
+    assert len(obs.messages) == 1, "After reset, should only have system prompt"
+    assert (
+        obs.messages[0]["content"] == system_prompt
+    ), "System prompt should be preserved after reset"
+
+    print("✓ test_system_prompt_preserved passed")
+
+
+def test_token_history_accumulation():
+    """Test that tokens accumulate correctly in the observation."""
+    tokenizer = MockTokenizer()
+    env = ChatEnvironment(tokenizer=tokenizer)
+
+    obs = env.reset()
+    initial_token_count = obs.tokens.numel()
+
+    # Step with first message
+    message1 = {"role": "user", "content": "Hi"}
+    action1 = env.message_to_action(message1)
+    obs1 = env.step(action1)
+    token_count_1 = obs1.tokens.numel()
+
+    # Tokens should increase
+    assert token_count_1 > initial_token_count, "Token count should increase after step"
+
+    # Step with second message
+    message2 = {"role": "assistant", "content": "Hello there"}
+    action2 = env.message_to_action(message2)
+    obs2 = env.step(action2)
+    token_count_2 = obs2.tokens.numel()
+
+    # Tokens should continue to accumulate
+    assert (
+        token_count_2 > token_count_1
+    ), "Token count should keep increasing with more messages"
+
+    # Verify tokens are the concatenation of both messages
+    expected_tokens = torch.cat([action1.tokens.flatten(), action2.tokens.flatten()])
+    assert torch.equal(
+        obs2.tokens, expected_tokens
+    ), "Tokens should be concatenation of all actions"
+
+    print("✓ test_token_history_accumulation passed")
+
+
+def test_direct_token_action():
+    """Test creating actions directly from tokens."""
+    tokenizer = MockTokenizer()
+    env = ChatEnvironment(tokenizer=tokenizer)
+
+    env.reset()
+
+    # Create raw tokens
+    raw_tokens = torch.tensor([[72, 101, 108, 108, 111]])  # ASCII for "Hello"
+    action = ChatAction(tokens=raw_tokens)
+
+    # Step with raw tokens
+    obs = env.step(action)
+
+    # Verify message was added
+    assert len(obs.messages) == 1, "Should have one message"
+    assert obs.messages[0]["role"] == "assistant", "Should default to assistant role"
+
+    # Verify tokens match what we sent (flattened)
+    assert torch.equal(
+        obs.tokens, raw_tokens.flatten()
+    ), "Observation tokens should match input tokens"
+
+    print("✓ test_direct_token_action passed")
+
+
+def test_empty_tokens_validation():
+    """Test that empty tokens raise a ValueError."""
+    try:
+        action = ChatAction(tokens=torch.tensor([]))
+        assert False, "Should have raised ValueError for empty tokens"
+    except ValueError as e:
+        assert "empty" in str(e).lower(), "Error message should mention empty tokens"
+
+    print("✓ test_empty_tokens_validation passed")
+
+
+def test_message_validation():
+    """Test that invalid messages raise appropriate errors."""
+    tokenizer = MockTokenizer()
+    env = ChatEnvironment(tokenizer=tokenizer)
+
+    # Test missing 'role' key
+    try:
+        env.message_to_action({"content": "test"})  # type: ignore
+        assert False, "Should have raised error for missing 'role' key"
+    except (ValueError, KeyError):
+        pass
+
+    # Test missing 'content' key
+    try:
+        env.message_to_action({"role": "user"})  # type: ignore
+        assert False, "Should have raised error for missing 'content' key"
+    except (ValueError, KeyError):
+        pass
+
+    # Test None content
+    try:
+        env.message_to_action({"role": "user", "content": None})  # type: ignore
+        assert False, "Should have raised error for None content"
+    except ValueError:
+        pass
+
+    print("✓ test_message_validation passed")
+
+
+def test_reset_clears_history():
+    """Test that reset properly clears all message and token history."""
+    tokenizer = MockTokenizer()
+    env = ChatEnvironment(tokenizer=tokenizer, system_prompt="System message")
+
+    # Add some messages
+    obs1 = env.reset()
+    initial_messages = len(obs1.messages)
+
+    action = env.message_to_action({"role": "user", "content": "Test message"})
+    obs2 = env.step(action)
+
+    # Verify message was added
+    assert (
+        len(obs2.messages) > initial_messages
+    ), "Message should be added after step"
+
+    # Reset
+    obs3 = env.reset()
+
+    # Verify we're back to just the system prompt
+    assert (
+        len(obs3.messages) == initial_messages
+    ), "Reset should clear history back to initial state"
+    assert (
+        obs3.messages[0]["content"] == "System message"
+    ), "System prompt should be preserved"
+
+    print("✓ test_reset_clears_history passed")
+
+
+def main():
+    """Run all tests."""
+    print("\n" + "=" * 60)
+    print("ChatEnvironment Test Suite")
+    print("=" * 60 + "\n")
+
+    tests = [
+        test_tokenization_consistency,
+        test_message_content_preservation,
+        test_system_prompt_preserved,
+        test_token_history_accumulation,
+        test_direct_token_action,
+        test_empty_tokens_validation,
+        test_message_validation,
+        test_reset_clears_history,
+    ]
+
+    failed = []
+    for test in tests:
+        try:
+            test()
+        except AssertionError as e:
+            print(f"✗ {test.__name__} failed: {e}")
+            failed.append(test.__name__)
+        except Exception as e:
+            print(f"✗ {test.__name__} errored: {e}")
+            import traceback
+
+            traceback.print_exc()
+            failed.append(test.__name__)
+
+    print("\n" + "=" * 60)
+    if not failed:
+        print(f"✓ All {len(tests)} tests passed!")
+        print("=" * 60)
+        return 0
+    else:
+        print(f"✗ {len(failed)}/{len(tests)} tests failed:")
+        for name in failed:
+            print(f"  - {name}")
+        print("=" * 60)
+        return 1
+
+
+if __name__ == "__main__":
+    exit(main())