feat: unify training rendering and harden FrozenLake VL

.github/workflows/training-ci.yml

@@ -0,0 +1,96 @@
+name: Training CI
+
+on:
+  pull_request:
+    paths:
+      - "training/**"
+      - ".github/workflows/training-ci.yml"
+  push:
+    branches:
+      - main
+    paths:
+      - "training/**"
+      - ".github/workflows/training-ci.yml"
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: training-ci-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  unit:
+    name: Unit And Import Tests
+    runs-on: ubuntu-latest
+    timeout-minutes: 20
+    defaults:
+      run:
+        working-directory: training
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set Up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+          cache: "pip"
+          cache-dependency-path: training/pyproject.toml
+
+      - name: Install Dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -e '.[dev]'
+
+      - name: Run Unit Tests
+        run: |
+          pytest -q tests/unit tests/test_smoke_imports.py examples/frozen_lake/test_masking.py
+
+  qwen3-4b-smoke:
+    name: Qwen3-4B Smoke (${{ matrix.test_target }})
+    needs: unit
+    runs-on: ubuntu-latest
+    timeout-minutes: 180
+    strategy:
+      fail-fast: false
+      max-parallel: 1
+      matrix:
+        test_target:
+          - tests/smoke_test/test_sft_smoke.py
+          - tests/smoke_test/test_dpo_smoke.py
+          - tests/smoke_test/test_grpo_smoke.py
+    defaults:
+      run:
+        working-directory: training
+    env:
+      FIREWORKS_API_KEY: ${{ secrets.FIREWORKS_API_KEY }}
+      FIREWORKS_ACCOUNT_ID: ${{ secrets.FIREWORKS_ACCOUNT_ID }}
+      FIREWORKS_BASE_URL: ${{ vars.FIREWORKS_BASE_URL || 'https://dev.api.fireworks.ai' }}
+      FIREWORKS_INFERENCE_URL: ${{ vars.FIREWORKS_INFERENCE_URL || vars.FIREWORKS_BASE_URL || 'https://dev.api.fireworks.ai' }}
+      FIREWORKS_HOTLOAD_API_URL: ${{ vars.FIREWORKS_HOTLOAD_API_URL || vars.FIREWORKS_BASE_URL || 'https://dev.api.fireworks.ai' }}
+      FIREWORKS_GATEWAY_SECRET: ${{ secrets.FIREWORKS_GATEWAY_SECRET }}
+      FIREWORKS_CUSTOM_IMAGE_TAG: ${{ vars.FIREWORKS_CUSTOM_IMAGE_TAG }}
+      FIREWORKS_SMOKE_BASE_MODEL: ${{ vars.FIREWORKS_SMOKE_BASE_MODEL || 'accounts/fireworks/models/qwen3-4b' }}
+      FIREWORKS_SMOKE_TOKENIZER_MODEL: ${{ vars.FIREWORKS_SMOKE_TOKENIZER_MODEL || 'Qwen/Qwen3-4B' }}
+      FIREWORKS_SMOKE_TRAINING_SHAPE: ${{ vars.FIREWORKS_SMOKE_TRAINING_SHAPE || 'ts-qwen3-4b-smoke-v1' }}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+
+      - name: Set Up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+          cache: "pip"
+          cache-dependency-path: training/pyproject.toml
+
+      - name: Install Dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install -e '.[dev]'
+
+      - name: Run Smoke Test
+        run: |
+          pytest -q -s ${{ matrix.test_target }}

training/examples/frozen_lake/frozen_lake_env.py

@@ -8,6 +8,7 @@
 from typing import Any, Dict, Iterable, List, Sequence, Tuple
 
 from training.examples.frozen_lake.frozen_lake_schema import FROZEN_LAKE_ACTIONS
+from training.examples.frozen_lake.rendering import render_frozen_lake_png_data_url
 
 ActionToDelta = {
     "LEFT": (0, -1),
@@ -150,12 +151,14 @@ def __init__(self, map_rows: Sequence[str], max_steps: int = 30):
         self._step_count = 0
         self._terminated = False
         self._truncated = False
+        self._last_action: str | None = None
 
     def reset(self) -> Dict[str, Any]:
         self._position = 0
         self._step_count = 0
         self._terminated = False
         self._truncated = False
+        self._last_action = None
         return self._current_state(action="RESET", reward=0.0)
 
     def step(self, action: str) -> Dict[str, Any]:
@@ -164,6 +167,7 @@ def step(self, action: str) -> Dict[str, Any]:
             raise ValueError(f"Invalid action '{action}'. Expected one of {FROZEN_LAKE_ACTIONS}")
 
         if self._terminated or self._truncated:
+            self._last_action = normalized_action
             return self._current_state(action=normalized_action, reward=0.0)
 
         row = self._position // self._side
@@ -174,6 +178,7 @@ def step(self, action: str) -> Dict[str, Any]:
 
         self._position = row * self._side + col
         self._step_count += 1
+        self._last_action = normalized_action
 
         tile = self._map_rows[row][col]
         reward = 1.0 if tile == "G" else 0.0
@@ -199,6 +204,17 @@ def _current_state(self, action: str, reward: float) -> Dict[str, Any]:
         )
         return result.as_tool_result()
 
+    def render_image_data_url(self, *, cell_size: int = 96) -> str:
+        row = self._position // self._side
+        col = self._position % self._side
+        return render_frozen_lake_png_data_url(
+            self._map_rows,
+            agent_row=row,
+            agent_col=col,
+            last_action=self._last_action,
+            cell_size=cell_size,
+        )
+
 
 def build_frozen_lake_tool_env(environment_context: Dict[str, Any] | None, max_steps: int) -> FrozenLakeToolEnv:
     """Create a FrozenLakeToolEnv from dataset ``environment_context``."""