Teacherenv

.secrets.baseline

@@ -133,7 +133,7 @@
         "filename": "README.md",
         "hashed_secret": "a8253456364f1bfc7da7ae4a1db5b45d106317a5",
         "is_verified": false,
-        "line_number": 454
+        "line_number": 530
       }
     ],
     "SLURM.md": [
@@ -561,5 +561,5 @@
       }
     ]
   },
-  "generated_at": "2026-03-02T22:46:56Z"
+  "generated_at": "2026-03-14T00:43:09Z"
 }

README.md

@@ -298,6 +298,82 @@ curl -s http://localhost:8002/latest_example | jq '{has_ids:(.distill_token_ids!
 - Trainers should validate alignment assumptions they require (sequence length, per-position top-k, etc.).
 - Teacher-side architecture and prompt/rendering strategy are intentionally out of scope for this PR.
 
+### TeacherDistillationEnv follow-up
+
+The follow-up teacher environment uses a dedicated teacher server config and
+attaches teacher prompt logprobs before the group is sent to the API.
+
+Teacher config shape:
+
+```python
+TeacherDistillationConfig(
+    teacher_enabled=True,
+    teacher_top_k=8,
+)
+```
+
+Teacher server configs are passed separately at init, just like the primary
+`server_configs`:
+
+```python
+env = MyTeacherEnv(
+    config=env_config,
+    server_configs=student_server_configs,
+    teacher_server_configs=[
+        APIServerConfig(
+            base_url="http://localhost:9003/v1",
+            model_name="Qwen/Qwen3-30B-A3B-Instruct-2507",
+            api_key="",
+            server_type="vllm",
+            tokenizer_name="Qwen/Qwen3-30B-A3B-Instruct-2507",
+        )
+    ],
+)
+```
+
+You can either:
+
+- build a teacher-enabled env by mixing `TeacherDistillationEnv` into an existing
+  `BaseEnv`-derived env such as `GSM8kEnv`, or
+- subclass `TeacherDistillationEnv` directly and implement the usual environment
+  methods yourself.
+
+In both cases, `TeacherDistillationEnv` still assumes the normal `BaseEnv`
+runtime contract: tokenized rollouts, `ScoredDataGroup` payloads, and the
+standard `handle_send_to_api(...)` transport path.
+
+CLI shape:
+
+```bash
+--env.teacher_enabled true \
+--teacher.base_url "http://localhost:9003/v1" \
+--teacher.model_name "Qwen/Qwen3-30B-A3B-Instruct-2507" \
+--teacher.server_type vllm \
+--env.teacher_top_k 8
+```
+
+If `--teacher.model_name` is a deployment alias rather than a tokenizer
+identifier, also set `--teacher.tokenizer_name ...` so the env can validate
+tokenizer compatibility.
+
+Scope note:
+
+- The teacher-aware CLI wiring currently exists for `serve`.
+- If `teacher_enabled=True`, the generic `process` and `evaluate` commands will
+  fail loudly at env construction time unless you instantiate the env yourself
+  and pass `teacher_server_configs=...`.
+
+Tokenizer requirement:
+
+- Teacher distillation currently requires the teacher and student to use the same tokenizer vocabulary.
+- If the tokenizers do not match, `TeacherDistillationEnv` raises an error instead of attempting token conversion.
+
+Why same-tokenizer is required:
+
+- `distill_token_ids` are consumed as student-vocabulary IDs by the trainer.
+- If the teacher uses a different vocabulary, the same integer token ID refers to different text on the teacher and student sides.
+- A decode/re-tokenize/remap pipeline is not a safe drop-in fix because it changes both token positions and token identities, which breaks the exact per-position token supervision that the current distillation loss assumes.
+
 ---
 
 ## Testing and Debugging Tools

atroposlib/envs/server_handling/openai_server.py

@@ -199,18 +199,14 @@ def resolve_openai_configs(
                 f"Error parsing multi-server OpenAI configuration from YAML under '{OPENAI_NAMESPACE}': {e}"
             ) from e
     elif isinstance(default_server_configs, APIServerConfig):
-        # Check APIServerConfig BEFORE ServerBaseline since APIServerConfig inherits from ServerBaseline
-        logger.info(
-            "Using single OpenAI server configuration based on merged settings (default/YAML/CLI)."
-        )
+        logger.info("Using single OpenAI server configuration.")
         try:
             final_openai_config = APIServerConfig(**openai_config_dict)
         except Exception as e:
             raise FailedExecutionException(
-                f"Error creating final OpenAI configuration from merged settings: {e}\n"
-                f"Merged Dict: {openai_config_dict}"
+                f"Error creating final OpenAI configuration: {e}"
             ) from e
-        server_configs = final_openai_config
+        server_configs = [final_openai_config]
     elif isinstance(default_server_configs, ServerBaseline):
         # Pure ServerBaseline (not APIServerConfig) - no CLI overrides possible
         logger.info("Using ServerBaseline configuration.")
@@ -219,26 +215,22 @@ def resolve_openai_configs(
         logger.info("Using default multi-server configuration (length >= 2).")
         server_configs = default_server_configs
     else:
-        logger.info(
-            "Using single OpenAI server configuration based on merged settings (default/YAML/CLI)."
-        )
+        logger.info("Using single OpenAI server configuration.")
         try:
             final_openai_config = APIServerConfig(**openai_config_dict)
         except Exception as e:
             raise FailedExecutionException(
-                f"Error creating final OpenAI configuration from merged settings: {e}\n"
-                f"Merged Dict: {openai_config_dict}"
+                f"Error creating final OpenAI configuration: {e}"
             ) from e
 
         if isinstance(default_server_configs, APIServerConfig):
-            server_configs = final_openai_config
+            server_configs = [final_openai_config]
         elif isinstance(default_server_configs, list):
             server_configs = [final_openai_config]
         else:
             logger.warning(
                 f"Unexpected type for default_server_configs: {type(default_server_configs)}. "
-                f"Proceeding with single OpenAI server configuration based on merged settings."
+                "Proceeding with single OpenAI server configuration."
             )
             server_configs = [final_openai_config]
-
     return server_configs

atroposlib/envs/server_handling/vllm_server.py

@@ -281,7 +281,7 @@ async def _get_logprobs_wrapper(self, **kwargs) -> Dict[str, Any]:
         ), "Prompt or input_ids is required for get_logprobs!"
 
         top_k = int(kwargs.pop("top_k", kwargs.pop("top_logprobs", 1)))
-        top_k = max(1, top_k)
+        top_k = max(0, top_k)
 
         # Use input_ids if provided (from ManagedServer), otherwise tokenize prompt
         from_prompt_text = False
@@ -408,25 +408,22 @@ def resolve_openai_configs(
         logger.info("Using default multi-server configuration (length >= 2).")
         server_configs = default_server_configs
     else:
-        logger.info(
-            "Using single OpenAI server configuration based on merged settings (default/YAML/CLI)."
-        )
+        logger.info("Using single OpenAI server configuration.")
         try:
             final_openai_config = APIServerConfig(**openai_config_dict)
         except Exception as e:
             raise FailedExecutionException(
-                f"Error creating final OpenAI configuration from merged settings: {e}\n"
-                f"Merged Dict: {openai_config_dict}"
+                f"Error creating final OpenAI configuration: {e}"
             ) from e
 
         if isinstance(default_server_configs, APIServerConfig):
-            server_configs = final_openai_config
+            server_configs = [final_openai_config]
         elif isinstance(default_server_configs, list):
             server_configs = [final_openai_config]
         else:
             logger.warning(
                 f"Unexpected type for default_server_configs: {type(default_server_configs)}. "
-                f"Proceeding with single OpenAI server configuration based on merged settings."
+                "Proceeding with single OpenAI server configuration."
             )
             server_configs = [final_openai_config]