File size: 4,145 Bytes
06c11b0 | 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 | """
test_TaskGoalIsList.py
Directly create a real Gymnasium environment (wrapped in DemonstrationWrapper),
call env.reset(), and verify that info["task_goal"] is a list and not empty.
Covers all 16 envs.
Run:
uv run python -m pytest tests/lightweight/test_TaskGoalI_isList.py -v -s
"""
import gymnasium as gym
import pytest
from typing import Literal
from robomme.env_record_wrapper.DemonstrationWrapper import DemonstrationWrapper
from robomme.env_record_wrapper.EndeffectorDemonstrationWrapper import EndeffectorDemonstrationWrapper
from robomme.env_record_wrapper.MultiStepDemonstrationWrapper import MultiStepDemonstrationWrapper
from robomme.env_record_wrapper.OraclePlannerDemonstrationWrapper import OraclePlannerDemonstrationWrapper
pytestmark = [pytest.mark.slow, pytest.mark.gpu]
# ── All 16 env_ids ─────────────────────────────────────────────────────────────
ALL_ENV_IDS = [
"BinFill",
"PickXtimes",
"SwingXtimes",
"StopCube",
"VideoUnmask",
"VideoUnmaskSwap",
"ButtonUnmask",
"ButtonUnmaskSwap",
"PickHighlight",
"VideoRepick",
"VideoPlaceButton",
"VideoPlaceOrder",
"MoveCube",
"InsertPeg",
"PatternLock",
"RouteStick",
]
# ── Four ActionSpaceTypes ──────────────────────────────────────────────────────
ACTION_SPACES = ["joint_angle", "ee_pose", "waypoint", "multi_choice"]
def _make_env(env_id: str, action_space: str):
"""Create and return a real environment wrapped with the corresponding Wrapper."""
env = gym.make(
env_id,
obs_mode="rgb+depth+segmentation",
control_mode="pd_joint_pos",
render_mode="rgb_array",
reward_mode="dense",
)
# Base Wrapper
env = DemonstrationWrapper(
env,
max_steps_without_demonstration=10002,
gui_render=False,
include_maniskill_obs=True,
include_front_depth=True,
include_wrist_depth=True,
include_front_camera_extrinsic=True,
include_wrist_camera_extrinsic=True,
include_available_multi_choices=True,
include_front_camera_intrinsic=True,
include_wrist_camera_intrinsic=True,
)
# Apply additional Wrapper depending on action_space
if action_space == "joint_angle":
pass
elif action_space == "ee_pose":
env = EndeffectorDemonstrationWrapper(env, action_repr="rpy")
elif action_space == "waypoint":
env = MultiStepDemonstrationWrapper(env, gui_render=False, vis=False)
elif action_space == "multi_choice":
env = OraclePlannerDemonstrationWrapper(env, env_id=env_id, gui_render=False)
else:
raise ValueError(f"Unsupported action_space: {action_space}")
return env
@pytest.mark.parametrize(
"env_id, action_space",
[(env, action) for env in ALL_ENV_IDS for action in ACTION_SPACES],
)
def test_task_goal_is_list(env_id: str, action_space: str):
"""
Test four action_spaces consecutively for each env_id:
1. Create real environment (including corresponding Wrapper)
2. Call reset()
3. Assert info["task_goal"] is list and not empty
"""
print(f"\nTesting [{env_id}] with action_space={action_space!r}")
env = _make_env(env_id, action_space)
try:
_, info = env.reset()
finally:
env.close()
task_goal = info["task_goal"]
print(f"[{env_id} | {action_space}] task_goal = {task_goal!r}")
assert isinstance(task_goal, list), (
f"[{env_id} | {action_space}] info['task_goal'] should be list, actually {type(task_goal).__name__!r}: {task_goal!r}"
)
assert len(task_goal) >= 1, (
f"[{env_id} | {action_space}] info['task_goal'] should not be empty list"
)
for i, item in enumerate(task_goal):
assert isinstance(item, str), (
f"[{env_id} | {action_space}] task_goal[{i}] should be str, actually {type(item).__name__!r}: {item!r}"
)
|