# Modality config for G1 BrainCo Picking tasks with EGO-ONLY camera
# Modified from g1_brainco_config.py
# CHANGE: HEAD CAMERA ONLY (no wrist cameras)

from gr00t.configs.data.embodiment_configs import register_modality_config
from gr00t.data.embodiment_tags import EmbodimentTag
from gr00t.data.types import (
    ActionConfig,
    ActionFormat,
    ActionRepresentation,
    ActionType,
    ModalityConfig,
)

g1_brainco_pick_ego_config = {
    # EGO-CENTRIC ONLY: single head camera (no wrist cameras)
    "video": ModalityConfig(
        delta_indices=[0],
        modality_keys=[
            "observation.images.head_stereo_left",  # HEAD ONLY (removed wrist cameras)
        ],
    ),
    # Proprioceptive state: arm joint positions + hand states
    "state": ModalityConfig(
        delta_indices=[0],
        modality_keys=[
            "observation.state.arm_q",      # 14D: arm joint positions
            "observation.state.hand_state", # 12D: BrainCo finger states
        ],
    ),
    # Actions: arm commands + hand commands
    "action": ModalityConfig(
        delta_indices=list(range(16)),  # 16-step horizon
        modality_keys=[
            "action.arm_q",       # 14D: arm joint targets (relative)
            "action.hand_cmd",    # 12D: hand commands (absolute)
        ],
        action_configs=[
            ActionConfig(
                rep=ActionRepresentation.RELATIVE,
                type=ActionType.NON_EEF,
                format=ActionFormat.DEFAULT,
                state_key="observation.state.arm_q",  # Relative to arm_q state
            ),
            ActionConfig(
                rep=ActionRepresentation.ABSOLUTE,
                type=ActionType.NON_EEF,
                format=ActionFormat.DEFAULT,
            ),
        ],
    ),
    # Language instruction (task description from episodes.jsonl)
    "language": ModalityConfig(
        delta_indices=[0],
        modality_keys=["annotation.human.task_description"],
    ),
}

register_modality_config(g1_brainco_pick_ego_config, embodiment_tag=EmbodimentTag.NEW_EMBODIMENT)