fix: preserve inner task names and raise on unknown task in EnvGroup

tests/test_env_group.py

@@ -353,8 +353,38 @@ def test_get_env_for_task(self, mock_client):
 
         assert env_group.get_env_for_task("math") == env1
         assert env_group.get_env_for_task("code") == env2
-        # Unknown task returns first environment as fallback
-        assert env_group.get_env_for_task("unknown") == env1
+        # Unknown task should raise rather than silently misroute
+        with pytest.raises(ValueError, match="No environment found for task"):
+            env_group.get_env_for_task("unknown")
+
+    def test_nested_env_group_preserves_inner_tasks(self, mock_client):
+        """Wrapping an EnvGroup in another EnvGroup must preserve inner task names."""
+        env1 = SingleTurnEnv(
+            client=mock_client,
+            model="test-model",
+            dataset=Dataset.from_dict({"question": ["q1"], "answer": ["a1"]}),
+            rubric=Rubric(),
+        )
+        env2 = SingleTurnEnv(
+            client=mock_client,
+            model="test-model",
+            dataset=Dataset.from_dict({"question": ["q2"], "answer": ["a2"]}),
+            rubric=Rubric(),
+        )
+
+        inner_group = EnvGroup(envs=[env1, env2], env_names=["math", "code"])
+        outer_group = EnvGroup(envs=[inner_group], env_names=["my_env"])
+
+        # Inner task names should be present in the outer env_map
+        assert outer_group.get_env_for_task("math") is inner_group
+        assert outer_group.get_env_for_task("code") is inner_group
+
+        # Dataset should retain the inner task labels
+        dataset = outer_group.get_dataset()
+        tasks = set(dataset["task"])
+        assert "math" in tasks
+        assert "code" in tasks
+        assert "my_env" not in tasks
 
     @pytest.mark.asyncio
     async def test_env_group_generate(self, mock_client, make_input):

verifiers/envs/env_group.py

@@ -179,18 +179,28 @@ def add_task(example):
             # Build dataset if using DatasetBuilder, returns None if not available
             env_dataset = env.build_dataset()
             if env_dataset is not None:
-                # override task column to use env_name for routing
-                if "task" in env_dataset.column_names:
-                    env_dataset = env_dataset.remove_columns(["task"])
-                env_dataset = env_dataset.map(add_task, **map_kwargs)
+                if isinstance(env, EnvGroup):
+                    # Preserve inner task names so routing works through both levels.
+                    # Register each inner task name pointing to the inner EnvGroup.
+                    for inner_name in env.env_map:
+                        self.env_map[inner_name] = env
+                else:
+                    # override task column to use env_name for routing
+                    if "task" in env_dataset.column_names:
+                        env_dataset = env_dataset.remove_columns(["task"])
+                    env_dataset = env_dataset.map(add_task, **map_kwargs)
                 datasets.append(env_dataset)
             # Build eval_dataset if using DatasetBuilder, returns None if not available
             env_eval_dataset = env.build_eval_dataset()
             if env_eval_dataset is not None:
-                # override task column to use env_name for routing
-                if "task" in env_eval_dataset.column_names:
-                    env_eval_dataset = env_eval_dataset.remove_columns(["task"])
-                env_eval_dataset = env_eval_dataset.map(add_task, **map_kwargs)
+                if isinstance(env, EnvGroup):
+                    for inner_name in env.env_map:
+                        self.env_map[inner_name] = env
+                else:
+                    # override task column to use env_name for routing
+                    if "task" in env_eval_dataset.column_names:
+                        env_eval_dataset = env_eval_dataset.remove_columns(["task"])
+                    env_eval_dataset = env_eval_dataset.map(add_task, **map_kwargs)
                 eval_datasets.append(env_eval_dataset)
         dataset = concatenate_datasets(datasets) if datasets else None
         eval_dataset = concatenate_datasets(eval_datasets) if eval_datasets else None
@@ -320,7 +330,13 @@ async def rollout(
         return await env.rollout(input, client, model, sampling_args)
 
     def get_env_for_task(self, task: str) -> vf.Environment:
-        return self.env_map.get(task, self.envs[0])
+        env = self.env_map.get(task)
+        if env is None:
+            available = list(self.env_map.keys())
+            raise ValueError(
+                f"No environment found for task '{task}'. Available tasks: {available}"
+            )
+        return env
 
     def set_max_seq_len(self, max_seq_len: int | None) -> None:
         """Set the max_seq_len value for this environment group and all sub-environments."""