workspace/tests/test_effect_rewards.py

from __future__ import annotations

import pytest

from dovla_cil.data.schema import ActionChunk
from dovla_cil.effects.extractors import extract_structured_effect
from dovla_cil.effects.failure_classifier import classify_failure, classify_toy_failure
from dovla_cil.effects.rewards import best_action_index, compute_reward, distance_to_target_reward
from dovla_cil.tasks.library import ToyTaskLibrary


def test_distance_reward_prefers_near_actions() -> None:
    assert distance_to_target_reward(0.1, success=False) > distance_to_target_reward(
        0.5, success=False
    )
    assert distance_to_target_reward(0.01, success=True) > 0.0


def test_distance_reward_rejects_negative() -> None:
    with pytest.raises(ValueError):
        distance_to_target_reward(-1.0, success=False)


def test_success_predicate_gives_success_reward() -> None:
    task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl")
    before = _state(red_mug=[0.0, 0.0, 0.03], blue_bowl=[1.0, 0.0, 0.03])
    after = _state(
        red_mug=[1.0, 0.0, 0.04],
        blue_bowl=[1.0, 0.0, 0.03],
        red_mug_extra={"inside": "blue_bowl"},
    )

    effect = extract_structured_effect(before, after, task=task)
    reward = compute_reward(task, effect)

    assert reward.success is True
    assert reward.terminal_success is True
    assert reward.progress == 1.0


def test_wrong_target_classified() -> None:
    task = ToyTaskLibrary().get_by_id("toy_pick_object_among_distractors")
    before = _state(red_mug=[0, 0, 0.03], blue_mug=[0.5, 0, 0.03], green_bowl=[1, 0, 0.03])
    after = _state(red_mug=[0, 0, 0.03], blue_mug=[0.8, 0, 0.03], green_bowl=[1, 0, 0.03])
    effect = extract_structured_effect(before, after, task=task)
    reward = compute_reward(task, effect)
    action = ActionChunk(
        representation="semantic",
        values=[{"command": "grasp", "object": "blue_mug"}],
        skill_type="grasp",
        metadata={"candidate_type": "wrong_target", "intended_target": "blue_mug"},
    )

    failure = classify_failure(task, action, effect, reward)

    assert failure.type == "wrong_target"


def test_no_motion_classified() -> None:
    task = ToyTaskLibrary().get_by_id("toy_lift_can")
    before = _state(can=[0, 0, 0.03])
    after = _state(can=[0, 0, 0.03])
    effect = extract_structured_effect(before, after, task=task)
    reward = compute_reward(task, effect)
    action = ActionChunk(
        representation="semantic",
        values=[{"command": "noop"}],
        skill_type="noop",
        metadata={"candidate_type": "noop", "intended_target": "can"},
    )

    failure = classify_failure(task, action, effect, reward)

    assert failure.type == "no_motion"


def test_partial_pick_place_classified() -> None:
    task = ToyTaskLibrary().get_by_id("toy_put_red_mug_in_blue_bowl")
    before = _state(red_mug=[0.0, 0.0, 0.03], blue_bowl=[1.0, 0.0, 0.03])
    after = _state(
        red_mug=[0.5, 0.0, 0.2],
        blue_bowl=[1.0, 0.0, 0.03],
        red_mug_extra={"grasped": True, "lifted": True},
    )
    effect = extract_structured_effect(
        before, after, task=task, rollout_info={"grasp_success": True}
    )
    reward = compute_reward(task, effect)
    action = ActionChunk(
        representation="semantic",
        values=[{"command": "grasp", "object": "red_mug"}],
        skill_type="grasp",
        metadata={"candidate_type": "near_miss", "intended_target": "red_mug"},
    )

    failure = classify_failure(task, action, effect, reward)

    assert reward.success is False
    assert 0.0 < reward.progress < 1.0
    assert failure.type == "partial_success"


def test_progress_is_bounded() -> None:
    task = ToyTaskLibrary().get_by_id("toy_push_cube_to_target_zone")
    before = _state(cube=[0, 0, 0.03], target_zone=[10, 0, 0.0])
    after = _state(cube=[100, 0, 0.03], target_zone=[10, 0, 0.0])
    effect = extract_structured_effect(before, after, task=task)
    reward = compute_reward(task, effect)

    assert 0.0 <= reward.progress <= 1.0


def test_effect_extraction_compatibility_and_toy_failure() -> None:
    effect = extract_structured_effect(
        {"position": 0.0}, {"position": 0.5}, info={"success": False}
    )
    assert effect.metrics["delta_position"] == 0.5
    assert effect.relation_after["success"] is False
    assert classify_toy_failure(distance=0.2, tolerance=0.05) == "missed_target"
    assert classify_toy_failure(distance=0.01, tolerance=0.05) is None


def test_best_action_index() -> None:
    assert best_action_index([0.0, 2.0, 1.0]) == 1


def _state(
    *,
    red_mug=None,
    blue_bowl=None,
    blue_mug=None,
    green_bowl=None,
    can=None,
    cube=None,
    target_zone=None,
    red_mug_extra=None,
) -> dict:
    objects = {}
    for name, position in {
        "red_mug": red_mug,
        "blue_bowl": blue_bowl,
        "blue_mug": blue_mug,
        "green_bowl": green_bowl,
        "can": can,
        "cube": cube,
        "target_zone": target_zone,
    }.items():
        if position is not None:
            objects[name] = {"position": list(position), "grasped": False, "lifted": False}
    if red_mug_extra and "red_mug" in objects:
        objects["red_mug"].update(red_mug_extra)
    return {"objects": objects, "near_threshold": 0.25, "lifted_z": 0.1}