diff --git a/scripts/visualize_task_env.py b/scripts/visualize_task_env.py index 8ed46c758..815f0bd2e 100644 --- a/scripts/visualize_task_env.py +++ b/scripts/visualize_task_env.py @@ -24,12 +24,11 @@ from __future__ import annotations import argparse -import dataclasses import sys import time from collections.abc import Sequence from pathlib import Path -from typing import TYPE_CHECKING, Any, cast, get_args, get_origin, get_type_hints +from typing import TYPE_CHECKING, Any, cast import numpy as np from unisim.backend.base import CameraCfg @@ -51,45 +50,6 @@ from unilab.base.scene import SceneCfg -def _build_reward_stub(env_cfg_cls: type) -> dict[str, Any] | None: - """Generate a minimal reward_config dict for the given env cfg class.""" - try: - type_hints = get_type_hints(env_cfg_cls) - except Exception: - return None - reward_type = type_hints.get("reward_config") - if reward_type is None: - return None - if get_origin(reward_type) is not None: - non_none = [a for a in get_args(reward_type) if a is not type(None)] - if not non_none: - return None - reward_type = non_none[0] - if not dataclasses.is_dataclass(reward_type): - return None - - stub: dict[str, Any] = {} - for f in dataclasses.fields(reward_type): - if f.default is not dataclasses.MISSING: - continue - if f.default_factory is not dataclasses.MISSING: # type: ignore[misc] - continue - type_str = str(f.type) - if "dict" in type_str: - stub[f.name] = {} - elif "list" in type_str or "tuple" in type_str or "Sequence" in type_str: - stub[f.name] = [] - elif "bool" in type_str: - stub[f.name] = False - elif "float" in type_str or "int" in type_str: - stub[f.name] = 0.0 - elif "str" in type_str: - stub[f.name] = "" - else: - stub[f.name] = None - return stub - - def _parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: parser = argparse.ArgumentParser( description="Visualize an RL task's training environment with zero actions." @@ -291,12 +251,7 @@ def _build_env_cfg_override(task_name: str) -> dict[str, Any]: f"Task '{task_name}' is not registered. " f"Available: {sorted(registry.list_registered_envs())}" ) - env_cfg = registry.materialize_env_config(task_name) - override: dict[str, Any] = {} - reward_stub = _build_reward_stub(type(env_cfg)) - if reward_stub is not None: - override["reward_config"] = reward_stub - return override + return {} def _print_backend_scene(env) -> None: diff --git a/src/unilab/utils/reward.py b/src/unilab/utils/reward.py index 679dd3015..7259ace99 100644 --- a/src/unilab/utils/reward.py +++ b/src/unilab/utils/reward.py @@ -19,9 +19,7 @@ def _to_reward_dict(value: object, *, error_message: str) -> RewardDict: return cast(RewardDict, resolved) -def resolve_reward_dict( - cfg: DictConfig, *, target_field: str = "reward_config" -) -> dict[str, RewardDict]: +def resolve_reward_dict(cfg: DictConfig, *, target_field: str = "rewards") -> dict[str, RewardDict]: """Extract and validate the reward config from a composed Hydra config. Args: @@ -33,12 +31,10 @@ def resolve_reward_dict( Raises: ValueError: If the reward config is missing or empty, or the - ``target_field`` is not a supported override key. + ``target_field`` is not the supported ``"rewards"`` override key. """ - if target_field not in {"reward_config", "rewards"}: - raise ValueError( - f"Reward config target_field must be 'reward_config' or 'rewards', got {target_field!r}" - ) + if target_field != "rewards": + raise ValueError(f"Reward config target_field must be 'rewards', got {target_field!r}") reward_cfg = OmegaConf.select(cfg, "reward") if not reward_cfg: raise ValueError( diff --git a/tests/config/test_reward_injection.py b/tests/config/test_reward_injection.py index 2268232ff..4ba2d438f 100644 --- a/tests/config/test_reward_injection.py +++ b/tests/config/test_reward_injection.py @@ -24,13 +24,29 @@ def test_resolve_reward_dict_reads_task_reward(): overrides=["task=go2_joystick_flat/mujoco"], ) - reward_dict = resolve_reward_dict(cfg)["reward_config"] + reward_dict = resolve_reward_dict(cfg)["rewards"] assert reward_dict["tracking_lin_vel"]["weight"] == 1.0 assert reward_dict["tracking_ang_vel"]["weight"] == 0.2 assert reward_dict["tracking_lin_vel"]["func"].endswith("track_lin_vel") +def test_resolve_reward_dict_rejects_reward_config_target(): + """The removed 'reward_config' override key must fail closed.""" + import pytest + + from unilab.utils.reward import resolve_reward_dict + + with initialize(config_path="../../src/unilab/conf/ppo", version_base="1.3"): + cfg = compose( + config_name="config", + overrides=["task=go2_joystick_flat/mujoco"], + ) + + with pytest.raises(ValueError, match="target_field must be 'rewards'"): + resolve_reward_dict(cfg, target_field="reward_config") + + def test_reward_config_conversion(): """Test reward config materializes into manager reward terms via registry.""" from unilab.base import registry