delinqu/gr00t-libero-goal

Gr00t-N1-libero-goal

from gr00t.eval.robot import RobotInferenceClient
from gr00t.experiment.data_config import DATA_CONFIG_MAP
from gr00t.model.policy import Gr00tPolicy

data_config = DATA_CONFIG_MAP[args.task_suite_name]
modality_config = data_config.modality_config()
transforms = data_config.transform()

policy_client = Gr00tPolicy(
        # host=args.host, 
        # port=args.port,
        model_path=args.pretrained_path,
        modality_config=modality_config,
        modality_transform=transforms,
        embodiment_tag="new_embodiment",
        device="cuda",
        denoising_steps=16
    )


batch = {
    "video.image": img[None],
    "video.wrist_image": wrist_img[None],
    "state.x": state[0:1][None],
    "state.y": state[1:2][None],
    "state.z": state[2:3][None],
    "state.roll": state[3:4][None],
    "state.pitch": state[4:5][None],
    "state.yaw": state[5:6][None],
    "state.gripper": state[6:][None],
    "annotation.human.action.task_description": [str(task_description)],
}

actions = policy_client.get_action(batch)
action_chunk = np.stack([
    actions["action.x"],
    actions["action.y"],
    actions["action.z"],
    actions["action.roll"],
    actions["action.pitch"],
    actions["action.yaw"],
    actions["action.gripper"],
], axis=-1) # [B, 7]