File size: 5,287 Bytes
20c251e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
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}