File size: 5,327 Bytes
52d0e83 cdd95ef f97c939 cdd95ef 52d0e83 cdd95ef f97c939 52d0e83 | 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 | from __future__ import annotations
import numpy as np
import pytest
from cil.chart_features import (
CONTEXT_HASH_WIDTH,
OBSERVATION_EMBED_DIM,
OBJECT_LAYOUT_EMBED_DIM,
build_chart_feature,
chart_feature_dim,
)
def test_base_chart_feature_preserves_original_flattened_base_action() -> None:
base_action = np.arange(12, dtype=np.float32).reshape(3, 4)
feature = build_chart_feature(base_action, mode="base")
assert feature.dtype == np.float32
assert np.allclose(feature, base_action.reshape(-1))
assert chart_feature_dim(base_action, mode="base") == 12
def test_base_context_feature_is_deterministic_and_extends_base() -> None:
base_action = np.zeros((2, 7), dtype=np.float32)
metadata = {"task_id": "PickCube-v1", "instruction": "pick up the cube"}
first = build_chart_feature(base_action, metadata, mode="base_context")
second = build_chart_feature(base_action, metadata, mode="base_context")
assert np.allclose(first, second)
assert first.shape == (base_action.size + 2 * CONTEXT_HASH_WIDTH + 2,)
assert chart_feature_dim(base_action, mode="base_context") == first.shape[0]
def test_base_context_feature_does_not_read_outcome_or_hidden_fields() -> None:
base_action = np.ones((1, 4), dtype=np.float32)
shared = {"task_id": "StackCube-v1", "instruction": "stack the cube"}
metadata_a = {
**shared,
"scalar_utility": 1.0,
"delta_U": 0.8,
"label": "positive",
"outcome_vector": {"success": True},
"hidden_chart_oracle": 1.0,
}
metadata_b = {
**shared,
"scalar_utility": -1.0,
"delta_U": -0.4,
"label": "negative",
"outcome_vector": {"success": False},
"hidden_chart_oracle": 0.0,
}
assert np.allclose(
build_chart_feature(base_action, metadata_a, mode="base_context"),
build_chart_feature(base_action, metadata_b, mode="base_context"),
)
def test_base_context_obs_appends_precomputed_observation_embedding(tmp_path) -> None:
base_action = np.zeros((2, 7), dtype=np.float32)
embeddings = np.arange(OBSERVATION_EMBED_DIM * 2, dtype=np.float32).reshape(
2, OBSERVATION_EMBED_DIM
)
path = tmp_path / "obs_embeddings.npz"
np.savez_compressed(path, embeddings=embeddings)
metadata = {
"task_id": "PullCube-v1",
"instruction": "pull the cube",
"observation_embedding_path": f"{path.name}#embeddings/1",
"_chart_root": str(tmp_path),
"scalar_utility": -999.0,
"label": "negative",
}
feature = build_chart_feature(base_action, metadata, mode="base_context_obs")
assert feature.shape == (
base_action.size + 2 * CONTEXT_HASH_WIDTH + 2 + OBSERVATION_EMBED_DIM,
)
assert np.allclose(feature[-OBSERVATION_EMBED_DIM:], embeddings[1])
assert chart_feature_dim(base_action, mode="base_context_obs") == feature.shape[0]
def test_base_context_obj_appends_precomputed_object_embedding(tmp_path) -> None:
base_action = np.zeros((2, 7), dtype=np.float32)
embeddings = np.arange(OBJECT_LAYOUT_EMBED_DIM * 2, dtype=np.float32).reshape(
2, OBJECT_LAYOUT_EMBED_DIM
)
path = tmp_path / "object_embeddings.npz"
np.savez_compressed(path, embeddings=embeddings)
metadata = {
"task_id": "PickCube-v1",
"instruction": "pick the cube",
"object_embedding_path": f"{path.name}#embeddings/1",
"_chart_root": str(tmp_path),
"outcome_vector": {"success": False},
"label": "negative",
}
feature = build_chart_feature(base_action, metadata, mode="base_context_obj")
expected = base_action.size + 2 * CONTEXT_HASH_WIDTH + 2 + OBJECT_LAYOUT_EMBED_DIM
assert feature.shape == (expected,)
assert np.allclose(feature[-OBJECT_LAYOUT_EMBED_DIM:], embeddings[1])
assert chart_feature_dim(base_action, mode="base_context_obj") == expected
def test_base_context_obs_obj_appends_both_embeddings(tmp_path) -> None:
base_action = np.zeros((1, 4), dtype=np.float32)
obs = np.ones((1, OBSERVATION_EMBED_DIM), dtype=np.float32)
obj = np.full((1, OBJECT_LAYOUT_EMBED_DIM), 2.0, dtype=np.float32)
obs_path = tmp_path / "obs.npz"
obj_path = tmp_path / "obj.npz"
np.savez_compressed(obs_path, embeddings=obs)
np.savez_compressed(obj_path, embeddings=obj)
metadata = {
"task_id": "StackCube-v1",
"instruction": "stack the cube",
"observation_embedding_path": f"{obs_path.name}#embeddings/0",
"object_embedding_path": f"{obj_path.name}#embeddings/0",
"_chart_root": str(tmp_path),
}
feature = build_chart_feature(base_action, metadata, mode="base_context_obs_obj")
expected = (
base_action.size
+ 2 * CONTEXT_HASH_WIDTH
+ 2
+ OBSERVATION_EMBED_DIM
+ OBJECT_LAYOUT_EMBED_DIM
)
assert feature.shape == (expected,)
assert np.allclose(feature[-(OBSERVATION_EMBED_DIM + OBJECT_LAYOUT_EMBED_DIM):-OBJECT_LAYOUT_EMBED_DIM], obs[0])
assert np.allclose(feature[-OBJECT_LAYOUT_EMBED_DIM:], obj[0])
def test_unknown_chart_feature_mode_raises() -> None:
with pytest.raises(ValueError, match="unknown chart feature mode"):
build_chart_feature(np.zeros((1, 2), dtype=np.float32), mode="outcome")
|