Upload folder using huggingface_hub

Dockerfile

@@ -0,0 +1,81 @@
+# 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 using openenv-base
+# This Dockerfile is flexible and works for both:
+# - In-repo environments (with local OpenEnv sources)
+# - Standalone environments (with openenv from PyPI/Git)
+# The build script (openenv build) handles context detection and sets appropriate build args.
+
+ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
+FROM ${BASE_IMAGE} AS builder
+
+WORKDIR /app
+
+# Ensure git is available (required for installing dependencies from VCS)
+RUN apt-get update && \
+    apt-get install -y --no-install-recommends git && \
+    rm -rf /var/lib/apt/lists/*
+
+# Build argument to control whether we're building standalone or in-repo
+ARG BUILD_MODE=in-repo
+ARG ENV_NAME=my_env
+
+# Copy environment code (always at root of build context)
+COPY . /app/env
+
+# For in-repo builds, openenv is already vendored in the build context
+# For standalone builds, openenv will be installed via pyproject.toml
+WORKDIR /app/env
+
+# Ensure uv is available (for local builds where base image lacks it)
+RUN if ! command -v uv >/dev/null 2>&1; then \
+        curl -LsSf https://astral.sh/uv/install.sh | sh && \
+        mv /root/.local/bin/uv /usr/local/bin/uv && \
+        mv /root/.local/bin/uvx /usr/local/bin/uvx; \
+    fi
+    
+# Install dependencies using uv sync
+# If uv.lock exists, use it; otherwise resolve on the fly
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -f uv.lock ]; then \
+        uv sync --frozen --no-install-project --no-editable; \
+    else \
+        uv sync --no-install-project --no-editable; \
+    fi
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    if [ -f uv.lock ]; then \
+        uv sync --frozen --no-editable; \
+    else \
+        uv sync --no-editable; \
+    fi
+
+# Final runtime stage
+FROM ${BASE_IMAGE}
+
+WORKDIR /app
+
+# Copy the virtual environment from builder
+COPY --from=builder /app/env/.venv /app/.venv
+
+# Copy the environment code
+COPY --from=builder /app/env /app/env
+
+# Set PATH to use the virtual environment
+ENV PATH="/app/.venv/bin:$PATH"
+
+# Set PYTHONPATH so imports work correctly
+ENV PYTHONPATH="/app/env:$PYTHONPATH"
+
+# 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
+# The module path is constructed to work with the /app/env structure
+ENV ENABLE_WEB_INTERFACE=true
+CMD ["sh", "-c", "cd /app/env && uvicorn server.app:app --host 0.0.0.0 --port 8000"]

README.md

@@ -1,10 +1,255 @@
 ---
-title: My Env
-emoji: 🦀
-colorFrom: gray
-colorTo: green
+title: My Env Environment Server
+emoji: 🎴
+colorFrom: blue
+colorTo: pink
 sdk: docker
 pinned: false
+app_port: 8000
+base_path: /web
+tags:
+  - openenv
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# My Env Environment
+
+A simple test environment that echoes back messages. Perfect for testing the env APIs as well as demonstrating environment usage patterns.
+
+## Quick Start
+
+The simplest way to use the My Env environment is through the `MyEnv` class:
+
+```python
+from my_env import MyAction, MyEnv
+
+try:
+    # Create environment from Docker image
+    my_envenv = MyEnv.from_docker_image("my_env-env:latest")
+
+    # Reset
+    result = my_envenv.reset()
+    print(f"Reset: {result.observation.echoed_message}")
+
+    # Send multiple messages
+    messages = ["Hello, World!", "Testing echo", "Final message"]
+
+    for msg in messages:
+        result = my_envenv.step(MyAction(message=msg))
+        print(f"Sent: '{msg}'")
+        print(f"  → Echoed: '{result.observation.echoed_message}'")
+        print(f"  → Length: {result.observation.message_length}")
+        print(f"  → Reward: {result.reward}")
+
+finally:
+    # Always clean up
+    my_envenv.close()
+```
+
+That's it! The `MyEnv.from_docker_image()` method handles:
+- Starting the Docker container
+- Waiting for the server to be ready
+- Connecting to the environment
+- Container cleanup when you call `close()`
+
+## Building the Docker Image
+
+Before using the environment, you need to build the Docker image:
+
+```bash
+# From project root
+docker build -t my_env-env:latest -f server/Dockerfile .
+```
+
+## Deploying to Hugging Face Spaces
+
+You can easily deploy your OpenEnv environment to Hugging Face Spaces using the `openenv push` command:
+
+```bash
+# From the environment directory (where openenv.yaml is located)
+openenv push
+
+# Or specify options
+openenv push --namespace my-org --private
+```
+
+The `openenv push` command will:
+1. Validate that the directory is an OpenEnv environment (checks for `openenv.yaml`)
+2. Prepare a custom build for Hugging Face Docker space (enables web interface)
+3. Upload to Hugging Face (ensuring you're logged in)
+
+### Prerequisites
+
+- Authenticate with Hugging Face: The command will prompt for login if not already authenticated
+
+### Options
+
+- `--directory`, `-d`: Directory containing the OpenEnv environment (defaults to current directory)
+- `--repo-id`, `-r`: Repository ID in format 'username/repo-name' (defaults to 'username/env-name' from openenv.yaml)
+- `--base-image`, `-b`: Base Docker image to use (overrides Dockerfile FROM)
+- `--private`: Deploy the space as private (default: public)
+
+### Examples
+
+```bash
+# Push to your personal namespace (defaults to username/env-name from openenv.yaml)
+openenv push
+
+# Push to a specific repository
+openenv push --repo-id my-org/my-env
+
+# Push with a custom base image
+openenv push --base-image ghcr.io/meta-pytorch/openenv-base:latest
+
+# Push as a private space
+openenv push --private
+
+# Combine options
+openenv push --repo-id my-org/my-env --base-image custom-base:latest --private
+```
+
+After deployment, your space will be available at:
+`https://huggingface.co/spaces/<repo-id>`
+
+The deployed space includes:
+- **Web Interface** at `/web` - Interactive UI for exploring the environment
+- **API Documentation** at `/docs` - Full OpenAPI/Swagger interface
+- **Health Check** at `/health` - Container health monitoring
+- **WebSocket** at `/ws` - Persistent session endpoint for low-latency interactions
+
+## Environment Details
+
+### Action
+**MyAction**: Contains a single field
+- `message` (str) - The message to echo back
+
+### Observation
+**MyObservation**: Contains the echo response and metadata
+- `echoed_message` (str) - The message echoed back
+- `message_length` (int) - Length of the message
+- `reward` (float) - Reward based on message length (length × 0.1)
+- `done` (bool) - Always False for echo environment
+- `metadata` (dict) - Additional info like step count
+
+### Reward
+The reward is calculated as: `message_length × 0.1`
+- "Hi" → reward: 0.2
+- "Hello, World!" → reward: 1.3
+- Empty message → reward: 0.0
+
+## Advanced Usage
+
+### Connecting to an Existing Server
+
+If you already have a My Env environment server running, you can connect directly:
+
+```python
+from my_env import MyEnv
+
+# Connect to existing server
+my_envenv = MyEnv(base_url="<ENV_HTTP_URL_HERE>")
+
+# Use as normal
+result = my_envenv.reset()
+result = my_envenv.step(MyAction(message="Hello!"))
+```
+
+Note: When connecting to an existing server, `my_envenv.close()` will NOT stop the server.
+
+### Using the Context Manager
+
+The client supports context manager usage for automatic connection management:
+
+```python
+from my_env import MyAction, MyEnv
+
+# Connect with context manager (auto-connects and closes)
+with MyEnv(base_url="http://localhost:8000") as env:
+    result = env.reset()
+    print(f"Reset: {result.observation.echoed_message}")
+    # Multiple steps with low latency
+    for msg in ["Hello", "World", "!"]:
+        result = env.step(MyAction(message=msg))
+        print(f"Echoed: {result.observation.echoed_message}")
+```
+
+The client uses WebSocket connections for:
+- **Lower latency**: No HTTP connection overhead per request
+- **Persistent session**: Server maintains your environment state
+- **Efficient for episodes**: Better for many sequential steps
+
+### Concurrent WebSocket Sessions
+
+The server supports multiple concurrent WebSocket connections. To enable this,
+modify `server/app.py` to use factory mode:
+
+```python
+# In server/app.py - use factory mode for concurrent sessions
+app = create_app(
+    MyEnvironment,  # Pass class, not instance
+    MyAction,
+    MyObservation,
+    max_concurrent_envs=4,  # Allow 4 concurrent sessions
+)
+```
+
+Then multiple clients can connect simultaneously:
+
+```python
+from my_env import MyAction, MyEnv
+from concurrent.futures import ThreadPoolExecutor
+
+def run_episode(client_id: int):
+    with MyEnv(base_url="http://localhost:8000") as env:
+        result = env.reset()
+        for i in range(10):
+            result = env.step(MyAction(message=f"Client {client_id}, step {i}"))
+        return client_id, result.observation.message_length
+
+# Run 4 episodes concurrently
+with ThreadPoolExecutor(max_workers=4) as executor:
+    results = list(executor.map(run_episode, range(4)))
+```
+
+## Development & Testing
+
+### Direct Environment Testing
+
+Test the environment logic directly without starting the HTTP server:
+
+```bash
+# From the server directory
+python3 server/my_env_environment.py
+```
+
+This verifies that:
+- Environment resets correctly
+- Step executes actions properly
+- State tracking works
+- Rewards are calculated correctly
+
+### Running Locally
+
+Run the server locally for development:
+
+```bash
+uvicorn server.app:app --reload
+```
+
+## Project Structure
+
+```
+my_env/
+├── .dockerignore         # Docker build exclusions
+├── __init__.py            # Module exports
+├── README.md              # This file
+├── openenv.yaml           # OpenEnv manifest
+├── pyproject.toml         # Project metadata and dependencies
+├── uv.lock                # Locked dependencies (generated)
+├── client.py              # MyEnv client
+├── models.py              # Action and Observation models
+└── server/
+    ├── __init__.py        # Server module exports
+    ├── my_env_environment.py  # Core environment logic
+    ├── app.py             # FastAPI application (HTTP + WebSocket endpoints)
+    └── Dockerfile         # Container image definition
+```

__init__.py

@@ -0,0 +1,23 @@
+# 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.
+
+"""Meta-Optimizer and My Env environments."""
+
+from .client import MetaOptimizerEnv
+from .models import (
+    MetaOptimizerAction,
+    MetaOptimizerObservation,
+    MyAction,
+    MyObservation,
+)
+
+__all__ = [
+    "MetaOptimizerEnv",
+    "MetaOptimizerAction",
+    "MetaOptimizerObservation",
+    "MyAction",
+    "MyObservation",
+]

client.py

@@ -0,0 +1,59 @@
+# 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.
+
+"""Meta-Optimizer Environment Client (OpenEnv WebSocket client)."""
+
+from typing import Dict
+
+from openenv.core.client_types import StepResult
+from openenv.core.env_server.types import State
+from openenv.core import EnvClient
+
+from .models import MetaOptimizerAction, MetaOptimizerObservation
+
+
+class MetaOptimizerEnv(
+    EnvClient[MetaOptimizerAction, MetaOptimizerObservation, State]
+):
+    """
+    Client for the Meta-Optimizer Environment.
+
+    Connects to the meta-optimizer server over WebSocket. Use reset(seed=..., task_id=...)
+    for training (task_id=None samples from 50 train tasks) or eval (task_id in EVAL_TASK_IDS).
+    """
+
+    def _step_payload(self, action: MetaOptimizerAction) -> Dict:
+        return {
+            "lr_scale": action.lr_scale,
+            "momentum_coef": action.momentum_coef,
+            "grad_clip_threshold": action.grad_clip_threshold,
+            "weight_decay_this_step": action.weight_decay_this_step,
+        }
+
+    def _parse_result(
+        self, payload: Dict
+    ) -> StepResult[MetaOptimizerObservation]:
+        obs_data = payload.get("observation", {})
+        observation = MetaOptimizerObservation(
+            loss=obs_data.get("loss", 0.0),
+            step_count=obs_data.get("step_count", 0),
+            grad_norm=obs_data.get("grad_norm"),
+            steps_to_threshold=obs_data.get("steps_to_threshold"),
+            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) -> State:
+        return State(
+            episode_id=payload.get("episode_id"),
+            step_count=payload.get("step_count", 0),
+        )

env_gym.py

@@ -0,0 +1,108 @@
+# 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.
+
+"""
+Gymnasium wrapper for MetaOptimizerEnvironment for use with Stable-Baselines3 (e.g. SAC).
+"""
+
+import math
+from typing import Any, Dict, Optional, Tuple
+
+import gymnasium as gym
+import numpy as np
+
+from my_env.models import MetaOptimizerAction
+from my_env.server.meta_optimizer_environment import MetaOptimizerEnvironment
+from my_env.server.tasks import get_task
+
+# Bounds for normalization / clipping
+LOSS_LOG_MAX = 2.0   # log10(loss+1e-8) capped for obs
+GRAD_NORM_SCALE = 10.0
+
+
+def obs_to_vector(obs: Any, max_steps: int) -> np.ndarray:
+    """Convert MetaOptimizerObservation to a fixed-size vector for SAC."""
+    loss = getattr(obs, "loss", 0.0) or 0.0
+    step_count = getattr(obs, "step_count", 0)
+    grad_norm = getattr(obs, "grad_norm", None)
+    # Normalize: log loss (bounded), step ratio, grad norm scale
+    loss_feat = min(math.log10(loss + 1e-8), LOSS_LOG_MAX) / LOSS_LOG_MAX
+    step_feat = step_count / max(1, max_steps)
+    grad_feat = (grad_norm / GRAD_NORM_SCALE) if grad_norm is not None else 0.0
+    grad_feat = min(max(grad_feat, 0.0), 1.0)
+    return np.array([loss_feat, step_feat, grad_feat], dtype=np.float32)
+
+
+def vector_to_action(vec: np.ndarray) -> MetaOptimizerAction:
+    """Map [0,1]^4 to action bounds: lr [1e-4, 1], momentum [0,1], clip [0, 2], wd [0, 1e-3]."""
+    lr = 1e-4 + (1.0 - 1e-4) * float(np.clip(vec[0], 0, 1))
+    momentum = float(np.clip(vec[1], 0, 1))
+    clip = 2.0 * float(np.clip(vec[2], 0, 1))
+    wd = 1e-3 * float(np.clip(vec[3], 0, 1))
+    return MetaOptimizerAction(
+        lr_scale=lr,
+        momentum_coef=momentum,
+        grad_clip_threshold=clip,
+        weight_decay_this_step=wd,
+    )
+
+
+class MetaOptimizerGymEnv(gym.Env):
+    """
+    Gymnasium env wrapping MetaOptimizerEnvironment for SAC.
+    Samples tasks from Distribution A (task_id 0..49) on each reset.
+    """
+
+    def __init__(
+        self,
+        max_steps: int = 100,
+        loss_threshold: float = 0.1,
+        task_ids: Optional[list] = None,
+    ):
+        super().__init__()
+        self._max_steps = max_steps
+        self._loss_threshold = loss_threshold
+        self._task_ids = task_ids or list(range(50))
+        self._env = MetaOptimizerEnvironment(
+            max_steps=max_steps,
+            loss_threshold=loss_threshold,
+        )
+        # Obs: loss (norm), step (norm), grad_norm (norm) = 3
+        self.observation_space = gym.spaces.Box(
+            low=0.0, high=1.0, shape=(3,), dtype=np.float32
+        )
+        # Action: lr, momentum, grad_clip, weight_decay (all [0,1] mapped to bounds in vector_to_action)
+        self.action_space = gym.spaces.Box(
+            low=0.0, high=1.0, shape=(4,), dtype=np.float32
+        )
+
+    def reset(
+        self, *, seed: Optional[int] = None, options: Optional[Dict] = None
+    ) -> Tuple[np.ndarray, Dict]:
+        import random
+        if seed is not None:
+            self._np_random = np.random.default_rng(seed)
+            idx = self._np_random.integers(0, len(self._task_ids))
+            task_id = self._task_ids[idx]
+        else:
+            task_id = random.choice(self._task_ids)
+        obs = self._env.reset(seed=seed, task_id=task_id)
+        vec = obs_to_vector(obs, self._max_steps)
+        return vec, {"task_id": task_id}
+
+    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict]:
+        act = vector_to_action(action)
+        obs = self._env.step(act)
+        vec = obs_to_vector(obs, self._max_steps)
+        reward = float(obs.reward if obs.reward is not None else 0.0)
+        done = bool(obs.done)
+        truncated = False
+        info = {
+            "loss": obs.loss,
+            "step_count": obs.step_count,
+            "steps_to_threshold": obs.steps_to_threshold,
+        }
+        return vec, reward, done, truncated, info

models.py

@@ -0,0 +1,75 @@
+# 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 My Env Environment.
+
+The my_env environment is a simple test environment that echoes back messages.
+Meta-optimizer models support the meta-learning RL optimizer environment.
+"""
+
+from pydantic import Field
+
+from openenv.core.env_server.types import Action, Observation
+
+
+class MyAction(Action):
+    """Action for the My Env environment - just a message to echo."""
+
+    message: str = Field(..., description="Message to echo back")
+
+
+class MyObservation(Observation):
+    """Observation from the My Env environment - the echoed message."""
+
+    echoed_message: str = Field(default="", description="The echoed message")
+    message_length: int = Field(default=0, description="Length of the echoed message")
+
+
+# --- Meta-optimizer environment (meta-learning RL optimizer) ---
+
+
+class MetaOptimizerAction(Action):
+    """Action for the meta-optimizer environment: control optimizer hyperparameters per step."""
+
+    lr_scale: float = Field(
+        ...,
+        ge=1e-4,
+        le=1.0,
+        description="Learning rate scale for this step (e.g. 1e-4 to 1.0)",
+    )
+    momentum_coef: float = Field(
+        ...,
+        ge=0.0,
+        le=1.0,
+        description="Momentum coefficient (0 = no momentum, 1 = full carry)",
+    )
+    grad_clip_threshold: float = Field(
+        ...,
+        ge=0.0,
+        description="Gradient clipping threshold (0 = no clipping)",
+    )
+    weight_decay_this_step: float = Field(
+        ...,
+        ge=0.0,
+        description="Weight decay (L2) scale for this step (0 = no weight decay)",
+    )
+
+
+class MetaOptimizerObservation(Observation):
+    """Observation from the meta-optimizer environment: loss, step, and optional grad norm."""
+
+    loss: float = Field(..., description="Current loss after last update")
+    step_count: int = Field(..., description="Current step in the episode")
+    grad_norm: float | None = Field(
+        default=None,
+        description="Global gradient norm before last update (if available)",
+    )
+    steps_to_threshold: int | None = Field(
+        default=None,
+        description="Step at which loss first reached threshold (None if not yet reached)",
+    )
+    

openenv.yaml

@@ -0,0 +1,7 @@
+spec_version: 1
+name: my_env
+type: space
+runtime: fastapi
+app: server.app:app
+port: 8000
+

openenv_my_env.egg-info/PKG-INFO

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: openenv-my_env
+Version: 0.1.0
+Summary: My Env environment for OpenEnv
+Requires-Python: >=3.10
+Requires-Dist: openenv-core[core]>=0.2.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: matplotlib>=3.5.0
+Requires-Dist: seaborn>=0.12.0
+Requires-Dist: gymnasium>=0.29.0
+Requires-Dist: stable-baselines3>=2.0.0
+Requires-Dist: numpy>=1.20.0
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == "dev"
+Requires-Dist: pytest-cov>=4.0.0; extra == "dev"

openenv_my_env.egg-info/SOURCES.txt

@@ -0,0 +1,20 @@
+README.md
+__init__.py
+client.py
+models.py
+pyproject.toml
+./__init__.py
+./client.py
+./env_gym.py
+./models.py
+openenv_my_env.egg-info/PKG-INFO
+openenv_my_env.egg-info/SOURCES.txt
+openenv_my_env.egg-info/dependency_links.txt
+openenv_my_env.egg-info/entry_points.txt
+openenv_my_env.egg-info/requires.txt
+openenv_my_env.egg-info/top_level.txt
+server/__init__.py
+server/app.py
+server/meta_optimizer_environment.py
+server/my_env_environment.py
+server/tasks.py

openenv_my_env.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

openenv_my_env.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+server = my_env.server.app:main

openenv_my_env.egg-info/requires.txt

@@ -0,0 +1,11 @@
+openenv-core[core]>=0.2.0
+torch>=2.0.0
+matplotlib>=3.5.0
+seaborn>=0.12.0
+gymnasium>=0.29.0
+stable-baselines3>=2.0.0
+numpy>=1.20.0
+
+[dev]
+pytest>=8.0.0
+pytest-cov>=4.0.0

openenv_my_env.egg-info/top_level.txt

@@ -0,0 +1 @@
+my_env

pyproject.toml

@@ -0,0 +1,49 @@
+# 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.
+
+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "openenv-my_env"
+version = "0.1.0"
+description = "My Env environment for OpenEnv"
+requires-python = ">=3.10"
+dependencies = [
+    # Core OpenEnv runtime (provides FastAPI server + HTTP client types)
+    # install from github
+    # "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git",
+    "openenv-core[core]>=0.2.0",
+    # Environment-specific dependencies
+    # Add all dependencies needed for your environment here
+    # Examples:
+    # "numpy>=1.19.0",
+    "torch>=2.0.0",
+    "matplotlib>=3.5.0",
+    "seaborn>=0.12.0",
+    "gymnasium>=0.29.0",
+    "stable-baselines3>=2.0.0",
+    "numpy>=1.20.0",
+    # "openspiel>=1.0.0",
+    # "smolagents>=1.22.0,<2",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-cov>=4.0.0",
+]
+
+[project.scripts]
+# Server entry point - enables running via: uv run --project . server
+# or: python -m my_env.server.app
+server = "my_env.server.app:main"
+
+[tool.setuptools]
+include-package-data = true
+packages = ["my_env", "my_env.server"]
+package-dir = { "my_env" = ".", "my_env.server" = "server" }

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.
+
+"""My Env environment server components."""
+
+from .my_env_environment import MyEnvironment
+
+__all__ = ["MyEnvironment"]

server/app.py

@@ -0,0 +1,81 @@
+# 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 Meta-Optimizer Environment.
+
+This module creates an HTTP server that exposes the MetaOptimizerEnvironment
+over HTTP and WebSocket endpoints, compatible with EnvClient.
+
+Endpoints:
+    - POST /reset: Reset the environment (optionally with task_id for eval)
+    - POST /step: Execute an action (lr_scale, momentum_coef, grad_clip_threshold, weight_decay_this_step)
+    - GET /state: Get current environment state
+    - GET /schema: Get action/observation schemas
+    - WS /ws: WebSocket endpoint for persistent sessions
+
+Usage:
+    # Development (with auto-reload):
+    uvicorn server.app:app --reload --host 0.0.0.0 --port 8000
+
+    # Production:
+    uvicorn server.app:app --host 0.0.0.0 --port 8000 --workers 4
+
+    # Or run directly:
+    python -m server.app
+"""
+
+try:
+    from openenv.core.env_server.http_server import create_app
+except Exception as e:  # pragma: no cover
+    raise ImportError(
+        "openenv is required for the web interface. Install dependencies with '\n    uv sync\n'"
+    ) from e
+
+# Import from package so server works when run via uv run server (my_env.server.app)
+from my_env.models import MetaOptimizerAction, MetaOptimizerObservation
+from .meta_optimizer_environment import MetaOptimizerEnvironment
+
+
+# Create the app with web interface and README integration
+app = create_app(
+    MetaOptimizerEnvironment,
+    MetaOptimizerAction,
+    MetaOptimizerObservation,
+    env_name="meta_optimizer",
+    max_concurrent_envs=4,
+)
+
+
+def main(host: str = "0.0.0.0", port: int = 8000):
+    """
+    Entry point for direct execution via uv run or python -m.
+
+    This function enables running the server without Docker:
+        uv run --project . server
+        uv run --project . server --port 8001
+        python -m my_env.server.app
+
+    Args:
+        host: Host address to bind to (default: "0.0.0.0")
+        port: Port number to listen on (default: 8000)
+
+    For production deployments, consider using uvicorn directly with
+    multiple workers:
+        uvicorn my_env.server.app:app --workers 4
+    """
+    import uvicorn
+
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=8000)
+    args = parser.parse_args()
+    main(port=args.port)

server/meta_optimizer_environment.py

@@ -0,0 +1,351 @@
+# 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.
+
+"""
+Meta-optimizer environment: train an RL agent to act as an optimizer on random regression tasks.
+
+Supports 50 training tasks, held-out eval, rich action space (LR, momentum, grad clip, weight decay),
+and convergence-speed reward. Action log is exposed for emergent-behavior visualization.
+"""
+
+import math
+import random
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
+
+import torch
+import torch.nn as nn
+
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
+from my_env.models import MetaOptimizerAction, MetaOptimizerObservation
+from .tasks import TRAIN_TASK_IDS, get_task, task_spec_from_dict, TaskSpec
+
+# Defaults
+LOSS_THRESHOLD = 0.1
+MAX_STEPS = 100
+BATCH_SIZE = 32
+
+
+def _build_model(spec: TaskSpec) -> nn.Module:
+    """Build a 2-layer MLP for the given task spec."""
+    torch.manual_seed(spec.arch_seed)
+    return nn.Sequential(
+        nn.Linear(spec.input_dim, spec.hidden_dim),
+        nn.ReLU(),
+        nn.Linear(spec.hidden_dim, spec.output_dim),
+    )
+
+
+def _get_batch(spec: TaskSpec, step: int, device: torch.device):
+    """Sinusoidal regression: X in [0,1], y = amplitude * sin(2*pi*freq*x + phase) + noise."""
+    g = torch.Generator(device=device)
+    g.manual_seed(spec.data_seed + step)
+    X = torch.rand(BATCH_SIZE, spec.input_dim, device=device, generator=g)
+    # y = amplitude * sin(2*pi*freq*x + phase); x is first column
+    x = X[:, 0:1]
+    y = spec.amplitude * torch.sin(2 * math.pi * spec.freq * x + spec.phase)
+    y = y + 0.05 * torch.randn_like(y, device=device, generator=g)
+    return X, y
+
+
+def run_adam_baseline(
+    task_id: Optional[int] = None,
+    task_spec: Optional[Dict[str, Any]] = None,
+    max_steps: int = MAX_STEPS,
+    loss_threshold: float = LOSS_THRESHOLD,
+    lr: float = 1e-2,
+    seed: Optional[int] = None,
+    return_metrics: bool = False,
+):
+    """
+    Run Adam on one task. Returns steps to threshold, or full metrics dict if return_metrics=True.
+    """
+    if (task_id is None) == (task_spec is None):
+        raise ValueError("Provide exactly one of task_id or task_spec")
+    if seed is not None:
+        torch.manual_seed(seed)
+    device = torch.device("cpu")
+    spec = task_spec_from_dict(task_spec) if task_spec is not None else get_task(task_id)
+    model = _build_model(spec).to(device)
+    opt = torch.optim.Adam(model.parameters(), lr=lr)
+    loss_trajectory: List[float] = []
+    steps_to_threshold: Optional[int] = None
+    for step in range(max_steps):
+        X, y = _get_batch(spec, step, device)
+        model.train()
+        opt.zero_grad()
+        loss = nn.functional.mse_loss(model(X), y)
+        loss.backward()
+        opt.step()
+        with torch.no_grad():
+            L = nn.functional.mse_loss(model(X), y).item()
+        loss_trajectory.append(L)
+        if steps_to_threshold is None and L < loss_threshold:
+            steps_to_threshold = step + 1
+    final_loss = loss_trajectory[-1] if loss_trajectory else float("inf")
+    if not return_metrics:
+        return steps_to_threshold if steps_to_threshold is not None else max_steps
+    last_k = min(10, len(loss_trajectory))
+    mean_last_k = sum(loss_trajectory[-last_k:]) / last_k if loss_trajectory else final_loss
+    return {
+        "steps_to_threshold": steps_to_threshold if steps_to_threshold is not None else max_steps,
+        "success": steps_to_threshold is not None,
+        "final_loss": final_loss,
+        "mean_last_10_loss": mean_last_k,
+        "loss_auc": sum(loss_trajectory) / len(loss_trajectory) if loss_trajectory else final_loss,
+        "loss_trajectory": loss_trajectory,
+    }
+
+
+def run_sgd_baseline(
+    task_id: Optional[int] = None,
+    task_spec: Optional[Dict[str, Any]] = None,
+    max_steps: int = MAX_STEPS,
+    loss_threshold: float = LOSS_THRESHOLD,
+    lr: float = 1e-2,
+    momentum: float = 0.9,
+    seed: Optional[int] = None,
+    return_metrics: bool = False,
+):
+    """
+    Run SGD (with optional momentum) on one task. Returns steps to threshold, or full metrics dict if return_metrics=True.
+    """
+    if (task_id is None) == (task_spec is None):
+        raise ValueError("Provide exactly one of task_id or task_spec")
+    if seed is not None:
+        torch.manual_seed(seed)
+    device = torch.device("cpu")
+    spec = task_spec_from_dict(task_spec) if task_spec is not None else get_task(task_id)
+    model = _build_model(spec).to(device)
+    opt = torch.optim.SGD(model.parameters(), lr=lr, momentum=momentum)
+    loss_trajectory = []
+    steps_to_threshold = None
+    for step in range(max_steps):
+        X, y = _get_batch(spec, step, device)
+        model.train()
+        opt.zero_grad()
+        loss = nn.functional.mse_loss(model(X), y)
+        loss.backward()
+        opt.step()
+        with torch.no_grad():
+            L = nn.functional.mse_loss(model(X), y).item()
+        loss_trajectory.append(L)
+        if steps_to_threshold is None and L < loss_threshold:
+            steps_to_threshold = step + 1
+    final_loss = loss_trajectory[-1] if loss_trajectory else float("inf")
+    if not return_metrics:
+        return steps_to_threshold if steps_to_threshold is not None else max_steps
+    last_k = min(10, len(loss_trajectory))
+    mean_last_k = sum(loss_trajectory[-last_k:]) / last_k if loss_trajectory else final_loss
+    return {
+        "steps_to_threshold": steps_to_threshold if steps_to_threshold is not None else max_steps,
+        "success": steps_to_threshold is not None,
+        "final_loss": final_loss,
+        "mean_last_10_loss": mean_last_k,
+        "loss_auc": sum(loss_trajectory) / len(loss_trajectory) if loss_trajectory else final_loss,
+        "loss_trajectory": loss_trajectory,
+    }
+
+
+def run_meta_optimizer_trajectory(
+    task_id: Optional[int] = None,
+    task_spec: Optional[Dict[str, Any]] = None,
+    max_steps: int = MAX_STEPS,
+    loss_threshold: float = LOSS_THRESHOLD,
+    seed: Optional[int] = None,
+    policy_callable: Optional[Any] = None,
+) -> Dict[str, Any]:
+    """
+    Run the meta-optimizer env with a policy (obs -> MetaOptimizerAction) and return metrics dict.
+    If policy_callable is None, uses a fixed default policy.
+    """
+    if (task_id is None) == (task_spec is None):
+        raise ValueError("Provide exactly one of task_id or task_spec")
+    if seed is not None:
+        random.seed(seed)
+        torch.manual_seed(seed)
+    env = MetaOptimizerEnvironment(max_steps=max_steps, loss_threshold=loss_threshold)
+    obs = env.reset(seed=seed, task_id=task_id, task_spec=task_spec)
+    loss_trajectory: List[float] = [obs.loss]
+    if policy_callable is None:
+        def _default_policy(o):  # type: ignore
+            return MetaOptimizerAction(
+                lr_scale=0.02, momentum_coef=0.9,
+                grad_clip_threshold=1.0, weight_decay_this_step=0.0,
+            )
+        policy_callable = _default_policy
+    while not obs.done:
+        action = policy_callable(obs)
+        obs = env.step(action)
+        loss_trajectory.append(obs.loss)
+    final_loss = obs.loss
+    steps_to_threshold = obs.steps_to_threshold if obs.steps_to_threshold is not None else max_steps
+    last_k = min(10, len(loss_trajectory))
+    mean_last_k = sum(loss_trajectory[-last_k:]) / last_k
+    return {
+        "steps_to_threshold": steps_to_threshold,
+        "success": obs.steps_to_threshold is not None,
+        "final_loss": final_loss,
+        "mean_last_10_loss": mean_last_k,
+        "loss_auc": sum(loss_trajectory) / len(loss_trajectory),
+        "loss_trajectory": loss_trajectory,
+    }
+
+
+class MetaOptimizerEnvironment(Environment[MetaOptimizerAction, MetaOptimizerObservation, State]):
+    """
+    Meta-learning optimizer environment: agent chooses LR scale, momentum, grad clip, weight decay per step.
+    Reward = -steps_to_reach_threshold (convergence speed). Supports 50 train tasks and held-out eval.
+    """
+
+    SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
+    def __init__(
+        self,
+        loss_threshold: float = LOSS_THRESHOLD,
+        max_steps: int = MAX_STEPS,
+        **kwargs: Any,
+    ):
+        super().__init__(**kwargs)
+        self.loss_threshold = loss_threshold
+        self.max_steps = max_steps
+        self._device = torch.device("cpu")
+
+        # Episode state (set in reset)
+        self._task_spec: Optional[TaskSpec] = None
+        self._model: Optional[nn.Module] = None
+        self._velocities: Optional[List[torch.Tensor]] = None
+        self._step_count: int = 0
+        self._current_loss: float = 0.0
+        self._steps_to_threshold: Optional[int] = None
+        self._action_log: List[Dict[str, Any]] = []
+        self._episode_id: Optional[str] = None
+
+    def reset(
+        self,
+        seed: Optional[int] = None,
+        episode_id: Optional[str] = None,
+        task_id: Optional[int] = None,
+        task_spec: Optional[Dict[str, Any]] = None,
+        **kwargs: Any,
+    ) -> MetaOptimizerObservation:
+        if seed is not None:
+            random.seed(seed)
+            torch.manual_seed(seed)
+        if task_spec is not None:
+            self._task_spec = task_spec_from_dict(task_spec)
+        else:
+            tid = task_id if task_id is not None else random.choice(TRAIN_TASK_IDS)
+            self._task_spec = get_task(tid)
+        self._model = _build_model(self._task_spec).to(self._device)
+        self._velocities = [torch.zeros_like(p) for p in self._model.parameters()]
+        self._step_count = 0
+        self._steps_to_threshold = None
+        self._action_log = []
+        self._episode_id = episode_id or str(uuid4())
+
+        # Initial loss (no update yet)
+        X, y = _get_batch(self._task_spec, 0, self._device)
+        with torch.no_grad():
+            out = self._model(X)
+            self._current_loss = nn.functional.mse_loss(out, y).item()
+
+        return self._observation(reward=None, grad_norm=None)
+
+    def step(
+        self,
+        action: MetaOptimizerAction,
+        timeout_s: Optional[float] = None,
+        **kwargs: Any,
+    ) -> MetaOptimizerObservation:
+        assert self._model is not None and self._task_spec is not None
+        lr = action.lr_scale
+        momentum = action.momentum_coef
+        clip = action.grad_clip_threshold
+        wd = action.weight_decay_this_step
+
+        self._action_log.append({
+            "step": self._step_count,
+            "lr_scale": lr,
+            "momentum_coef": momentum,
+            "grad_clip_threshold": clip,
+            "weight_decay_this_step": wd,
+        })
+
+        X, y = _get_batch(self._task_spec, self._step_count + 1, self._device)
+        self._model.train()
+        out = self._model(X)
+        loss = nn.functional.mse_loss(out, y)
+        self._model.zero_grad()
+        loss.backward()
+
+        grads = [p.grad.clone() for p in self._model.parameters()]
+        grad_norm = sum(g.pow(2).sum() for g in grads).sqrt().item()
+
+        if clip > 0:
+            total_norm = sum(g.pow(2).sum() for g in grads).sqrt()
+            if total_norm > clip:
+                scale = clip / (total_norm + 1e-8)
+                grads = [g * scale for g in grads]
+
+        with torch.no_grad():
+            for i, p in enumerate(self._model.parameters()):
+                g = grads[i]
+                v = self._velocities[i]
+                v.mul_(momentum).add_(g)
+                p.sub_(v, alpha=lr)
+                if wd > 0:
+                    p.sub_(p, alpha=wd)
+
+        with torch.no_grad():
+            new_out = self._model(X)
+            self._current_loss = nn.functional.mse_loss(new_out, y).item()
+
+        self._step_count += 1
+        if self._steps_to_threshold is None and self._current_loss < self.loss_threshold:
+            self._steps_to_threshold = self._step_count
+
+        done = self._step_count >= self.max_steps
+        if done:
+            reward = -(self._steps_to_threshold if self._steps_to_threshold is not None else self.max_steps)
+        else:
+            reward = 0.0
+
+        return self._observation(reward=reward, grad_norm=grad_norm, done=done)
+
+    def _observation(
+        self,
+        reward: Optional[float] = None,
+        grad_norm: Optional[float] = None,
+        done: bool = False,
+    ) -> MetaOptimizerObservation:
+        meta: Dict[str, Any] = {}
+        if self._steps_to_threshold is not None:
+            meta["steps_to_threshold"] = self._steps_to_threshold
+        if done and self._action_log:
+            meta["action_log"] = self._action_log
+        return MetaOptimizerObservation(
+            loss=self._current_loss,
+            step_count=self._step_count,
+            grad_norm=grad_norm,
+            steps_to_threshold=self._steps_to_threshold,
+            done=done,
+            reward=reward,
+            metadata=meta,
+        )
+
+    @property
+    def state(self) -> State:
+        return State(
+            episode_id=self._episode_id,
+            step_count=self._step_count,
+        )
+
+    def get_episode_action_log(self) -> List[Dict[str, Any]]:
+        """Return the action log for the current episode (for in-process viz or eval)."""
+        return list(self._action_log)

server/my_env_environment.py

@@ -0,0 +1,101 @@
+# 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.
+
+"""
+My Env Environment Implementation.
+
+A simple test environment that echoes back messages sent to it.
+Perfect for testing HTTP server infrastructure.
+"""
+
+from uuid import uuid4
+
+from openenv.core.env_server.interfaces import Environment
+from openenv.core.env_server.types import State
+
+from my_env.models import MyAction, MyObservation
+
+
+class MyEnvironment(Environment):
+    """
+    A simple echo environment that echoes back messages.
+
+    This environment is designed for testing the HTTP server infrastructure.
+    It maintains minimal state and simply echoes back whatever message it receives.
+
+    Example:
+        >>> env = MyEnvironment()
+        >>> obs = env.reset()
+        >>> print(obs.echoed_message)  # "My Env environment ready!"
+        >>>
+        >>> obs = env.step(MyAction(message="Hello"))
+        >>> print(obs.echoed_message)  # "Hello"
+        >>> print(obs.message_length)  # 5
+    """
+
+    # Enable concurrent WebSocket sessions.
+    # Set to True if your environment isolates state between instances.
+    # When True, multiple WebSocket clients can connect simultaneously, each
+    # getting their own environment instance (when using factory mode in app.py).
+    SUPPORTS_CONCURRENT_SESSIONS: bool = True
+
+    def __init__(self):
+        """Initialize the my_env environment."""
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._reset_count = 0
+
+    def reset(self) -> MyObservation:
+        """
+        Reset the environment.
+
+        Returns:
+            MyObservation with a ready message
+        """
+        self._state = State(episode_id=str(uuid4()), step_count=0)
+        self._reset_count += 1
+
+        return MyObservation(
+            echoed_message="My Env environment ready!",
+            message_length=0,
+            done=False,
+            reward=0.0,
+        )
+
+    def step(self, action: MyAction) -> MyObservation:  # type: ignore[override]
+        """
+        Execute a step in the environment by echoing the message.
+
+        Args:
+            action: MyAction containing the message to echo
+
+        Returns:
+            MyObservation with the echoed message and its length
+        """
+        self._state.step_count += 1
+
+        message = action.message
+        length = len(message)
+
+        # Simple reward: longer messages get higher rewards
+        reward = length * 0.1
+
+        return MyObservation(
+            echoed_message=message,
+            message_length=length,
+            done=False,
+            reward=reward,
+            metadata={"original_message": message, "step": self._state.step_count},
+        )
+
+    @property
+    def state(self) -> State:
+        """
+        Get the current environment state.
+
+        Returns:
+            Current State with episode_id and step_count
+        """
+        return self._state

server/requirements.txt

@@ -0,0 +1,6 @@
+openenv[core]>=0.2.0
+fastapi>=0.115.0
+uvicorn>=0.24.0
+
+
+

server/tasks.py

@@ -0,0 +1,115 @@
+# 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.
+
+"""
+Task registry for meta-learning.
+
+Tasks can be from the internal registry (get_task(task_id)) or provided from outside
+via task_spec_from_dict() — the client sends the task definition to the environment.
+"""
+
+from dataclasses import dataclass
+from typing import Any, Dict, List
+
+import math
+
+# Distribution A: 50 training tasks (low-freq sinusoids)
+TRAIN_TASK_IDS: List[int] = list(range(50))
+
+# Distribution B: held-out eval tasks (high-freq sinusoids — different distribution)
+EVAL_TASK_IDS: List[int] = [50, 51]
+
+# Bounds for each distribution (freq, amplitude, phase)
+DIST_A_FREQ = (1.0, 3.0)
+DIST_A_AMP = (0.5, 2.0)
+DIST_B_FREQ = (4.0, 6.0)
+DIST_B_AMP = (0.3, 1.5)
+
+
+@dataclass
+class TaskSpec:
+    """Spec for one sinusoidal regression task."""
+
+    task_id: int
+    input_dim: int  # 1 for scalar sinusoid input
+    hidden_dim: int
+    output_dim: int
+    data_seed: int
+    arch_seed: int
+    # Sinusoidal target: y = amplitude * sin(2*pi*freq*x + phase)
+    amplitude: float
+    freq: float
+    phase: float
+    distribution: str  # "A" or "B"
+
+
+def get_task(task_id: int) -> TaskSpec:
+    """
+    Return the task spec for the given task_id.
+    Task IDs 0..49 = Distribution A (train), 50+ = Distribution B (eval).
+    """
+    if task_id < 0:
+        raise ValueError(f"task_id must be >= 0, got {task_id}")
+    r = task_id * 7919 + 1
+    data_seed = task_id * 31337
+    arch_seed = task_id * 131 + 7
+    hidden_dim = 32 + (r % 33)
+
+    if task_id < 50:
+        # Distribution A
+        f_lo, f_hi = DIST_A_FREQ
+        a_lo, a_hi = DIST_A_AMP
+        distribution = "A"
+    else:
+        # Distribution B
+        f_lo, f_hi = DIST_B_FREQ
+        a_lo, a_hi = DIST_B_AMP
+        distribution = "B"
+
+    # Deterministic but varied per task
+    freq = f_lo + (r % 1000) / 1000.0 * (f_hi - f_lo)
+    amplitude = a_lo + ((r * 3) % 1000) / 1000.0 * (a_hi - a_lo)
+    phase = ((r * 7) % 1000) / 1000.0 * 2 * math.pi
+
+    return TaskSpec(
+        task_id=task_id,
+        input_dim=1,
+        hidden_dim=hidden_dim,
+        output_dim=1,
+        data_seed=data_seed,
+        arch_seed=arch_seed,
+        amplitude=amplitude,
+        freq=freq,
+        phase=phase,
+        distribution=distribution,
+    )
+
+
+def task_spec_from_dict(d: Dict[str, Any]) -> TaskSpec:
+    """
+    Build a TaskSpec from an external dict (sent by the client).
+    The task is defined outside the env; we just parse it here.
+
+    Expected keys for type "sinusoid":
+      type="sinusoid", amplitude, freq, phase, data_seed (optional), arch_seed (optional),
+      input_dim (optional, default 1), hidden_dim (optional, default 32), task_id (optional).
+    """
+    task_type = d.get("type", "sinusoid")
+    if task_type != "sinusoid":
+        raise ValueError(f"Unknown task type: {task_type}")
+    task_id = d.get("task_id", 0)
+    return TaskSpec(
+        task_id=task_id,
+        input_dim=int(d.get("input_dim", 1)),
+        hidden_dim=int(d.get("hidden_dim", 32)),
+        output_dim=1,
+        data_seed=int(d.get("data_seed", task_id * 31337)),
+        arch_seed=int(d.get("arch_seed", task_id * 131 + 7)),
+        amplitude=float(d["amplitude"]),
+        freq=float(d["freq"]),
+        phase=float(d["phase"]),
+        distribution=d.get("distribution", "external"),
+    )