Upload dreamzero source code (no model checkpoints)

fbd9366 verified 2 months ago

11.6 kB

from typing import Any, Dict, List, Optional

import numpy as np
from pydantic import Field, PrivateAttr
import torch

from groot.vla.data.schema import DatasetMetadata, StateActionMetadata
from groot.vla.data.transform.base import InvertibleModalityTransform, ModalityTransform


class ConcatTransform(InvertibleModalityTransform):
    """
    Concatenate the keys according to specified order.
    """

    # -- We inherit from ModalityTransform, so we keep apply_to as well --
    apply_to: list[str] = Field(
        default_factory=list, description="Not used in this transform, kept for compatibility."
    )

    video_concat_order: list[str] = Field(
        ...,
        description="Concatenation order for each video modality. "
        "Format: ['video.ego_view_pad_res224_freq20', ...]",
    )

    state_concat_order: Optional[list[str]] = Field(
        default=None,
        description="Concatenation order for each state modality. "
        "Format: ['state.position', 'state.velocity', ...].",
    )

    action_concat_order: Optional[list[str]] = Field(
        default=None,
        description="Concatenation order for each action modality. "
        "Format: ['action.position', 'action.velocity', ...].",
    )

    action_dims: dict[str, int] = Field(
        default_factory=dict,
        description="The dimensions of the action keys.",
    )
    state_dims: dict[str, int] = Field(
        default_factory=dict,
        description="The dimensions of the state keys.",
    )

    action_dims_post_transform: dict[str, int] = Field(
        default_factory=dict,
        description="The new dimensions of the action keys after transform is applied.",
    )
    state_dims_post_transform: dict[str, int] = Field(
        default_factory=dict,
        description="The new dimensions of the state keys after transform is applied.",
    )
    # Store the transform pipeline to examine for dimension changes
    _transform_pipeline: List[ModalityTransform] = PrivateAttr(default_factory=list)

    def model_dump(self, *args, **kwargs):
        if kwargs.get("mode", "python") == "json":
            include = {
                "apply_to",
                "video_concat_order",
                "state_concat_order",
                "action_concat_order",
            }
        else:
            include = kwargs.pop("include", None)

        return super().model_dump(*args, include=include, **kwargs)

    def set_transform_pipeline(self, transforms: List[ModalityTransform]):
        """Set the transform pipeline so this transform can examine it for dimension changes."""
        self._transform_pipeline = transforms

    def _get_target_rotations_from_pipeline(self) -> Dict[str, str]:
        """Extract target_rotations from StateActionTransform instances in the pipeline."""
        target_rotations = {}
        for transform in self._transform_pipeline:
            if hasattr(transform, "target_rotations"):
                transform_target_rotations = getattr(transform, "target_rotations", {})
                if transform_target_rotations:
                    target_rotations.update(transform_target_rotations)
        return target_rotations

    def apply(self, data: Dict[str, Any]) -> Dict[str, Any]:
        grouped_keys = {}
        for key in data.keys():
            try:
                modality, _ = key.split(".")
            except:  # noqa: E722
                ### Handle language annotation special case
                if "annotation" in key:
                    modality = "language"
                else:
                    modality = "others"
            if modality not in grouped_keys:
                grouped_keys[modality] = []
            grouped_keys[modality].append(key)

        if "video" in grouped_keys:
            # Check if keys in video_concat_order, state_concat_order, action_concat_order are
            # ineed contained in the data. If not, then the keys are misspecified
            video_keys = grouped_keys["video"]
            assert self.video_concat_order is not None, f"{self.video_concat_order=}, {video_keys=}"
            assert all(
                item in video_keys for item in self.video_concat_order
            ), f"keys in video_concat_order are misspecified, \n{video_keys=}, \n{self.video_concat_order=}"

            # Process each video view
            unsqueezed_videos = []
            for video_key in self.video_concat_order:
                video_data = data.pop(video_key)
                unsqueezed_video = np.expand_dims(
                    video_data, axis=-4
                )  # [..., H, W, C] -> [..., 1, H, W, C]
                unsqueezed_videos.append(unsqueezed_video)
            # Concatenate along the new axis
            unsqueezed_video = np.concatenate(unsqueezed_videos, axis=-4)  # [..., V, H, W, C]

            # Video
            data["video"] = unsqueezed_video

        # "state"
        if "state" in grouped_keys:
            state_keys = grouped_keys["state"]
            assert self.state_concat_order is not None, f"{self.state_concat_order=}"
            assert all(
                item in state_keys for item in self.state_concat_order
            ), f"keys in state_concat_order are misspecified, \n{state_keys=}, \n{self.state_concat_order=}"
            # Check the state dims
            for key in self.state_concat_order:
                target_shapes = [self.state_dims[key]]
                if self.is_rotation_key(key):
                    target_shapes.extend(
                        [3, 4, 6]
                    )  # 3 -> axis_angle, 4 -> quaternion, 6 -> rotation_6d
                target_shapes.append(self.state_dims[key] * 2)  # Allow for sin-cos transform
                assert (
                    data[key].shape[-1] in target_shapes
                ), f"State dim mismatch for {key=}, {data[key].shape[-1]=}, {target_shapes=}"
            # Concatenate the state keys
            # We'll have StateActionToTensor before this transform, so here we use torch.cat
            data["state"] = torch.cat(
                [data.pop(key) for key in self.state_concat_order], dim=-1
            )  # [T, D_state]

        if "action" in grouped_keys:
            action_keys = grouped_keys["action"]
            assert self.action_concat_order is not None, f"{self.action_concat_order=}"
            # Check if all keys in concat_order are present
            assert set(self.action_concat_order) == set(
                action_keys
            ), f"{set(self.action_concat_order)=}, {set(action_keys)=}"
            # Record the action dims
            for key in self.action_concat_order:
                target_shapes = [self.action_dims[key]]
                if self.is_rotation_key(key):
                    target_shapes.extend(
                        [3, 4, 6]
                    )  # 3 -> axis_angle, 4 -> quaternion, 6 -> rotation_6d
                assert (
                    data[key].shape[-1] in target_shapes
                ), f"Action dim mismatch for {key=}, {data[key].shape[-1]=}, {target_shapes=}"
            # Concatenate the action keys
            # We'll have StateActionToTensor before this transform, so here we use torch.cat
            data["action"] = torch.cat(
                [data.pop(key) for key in self.action_concat_order], dim=-1
            )  # [T, D_action]

        return data

    def unapply(self, data: dict) -> dict:
        start_dim = 0
        assert "action" in data, f"{data.keys()=}"
        # For those dataset without actions (LAPA), we'll never run unapply
        assert self.action_concat_order is not None, f"{self.action_concat_order=}"
        action_tensor = data.pop("action")
        for key in self.action_concat_order:
            if key not in self.action_dims:
                raise ValueError(f"Action dim {key} not found in action_dims.")
            end_dim = start_dim + self.get_state_action_dims_post_transform(key)
            data[key] = action_tensor[..., start_dim:end_dim]
            start_dim = end_dim
        if "state" in data:
            assert self.state_concat_order is not None, f"{self.state_concat_order=}"
            start_dim = 0
            state_tensor = data.pop("state")
            for key in self.state_concat_order:
                end_dim = start_dim + self.get_state_action_dims_post_transform(key)
                data[key] = state_tensor[..., start_dim:end_dim]
                start_dim = end_dim
        return data

    def __call__(self, data: dict) -> dict:
        return self.apply(data)

    def get_modality_metadata(self, key: str) -> StateActionMetadata:
        modality, subkey = key.split(".")
        assert self.dataset_metadata is not None, "Metadata not set"
        modality_config = getattr(self.dataset_metadata.modalities, modality)
        assert subkey in modality_config, f"{subkey=} not found in {modality_config=}"
        assert isinstance(
            modality_config[subkey], StateActionMetadata
        ), f"Expected {StateActionMetadata} for {subkey=}, got {type(modality_config[subkey])=}"
        return modality_config[subkey]

    def get_state_action_dims(self, key: str) -> int:
        """Get the dimension of a state or action key from the dataset metadata."""
        modality_config = self.get_modality_metadata(key)
        shape = modality_config.shape
        assert len(shape) == 1, f"{shape=}"
        return shape[0]

    def get_state_action_dims_post_transform(self, key: str) -> int:
        """
        This function is used to get the dims of the state/action keys after transform is applied.
        It is different from the `get_state_action_dims` function, because this function accounts for
        the case where we apply transforms and the # of dims is change eg. after applying axis_angle transform on
        quaternion, the dims change from 4D to 3D.
        """
        modality_config = self.get_modality_metadata(key)
        shape = modality_config.shape
        assert len(shape) == 1, f"{shape=}"

        if self.is_rotation_key(key):
            target_rotations = self._get_target_rotations_from_pipeline()
            if key in target_rotations:
                target_rotation = target_rotations[key]
                if target_rotation == "axis_angle":
                    return 3
                elif target_rotation == "quaternion":
                    return 4
                elif target_rotation == "rotation_6d":
                    return 6
                elif target_rotation == "euler_angles":
                    return 3
                else:
                    raise ValueError(f"Unknown target rotation type: {target_rotation}")
            else:
                # No target rotation specified, return original dimension
                return shape[0]
        else:
            return shape[0]

    def is_rotation_key(self, key: str) -> bool:
        modality_config = self.get_modality_metadata(key)
        return modality_config.rotation_type is not None

    def set_metadata(self, dataset_metadata: DatasetMetadata):
        """Set the metadata and compute the dimensions of the state and action keys."""
        super().set_metadata(dataset_metadata)
        # Pre-compute the dimensions of the state and action keys
        if self.action_concat_order is not None:
            for key in self.action_concat_order:
                self.action_dims[key] = self.get_state_action_dims(key)
        if self.state_concat_order is not None:
            for key in self.state_concat_order:
                self.state_dims[key] = self.get_state_action_dims(key)

	from typing import Any, Dict, List, Optional

	import numpy as np
	from pydantic import Field, PrivateAttr
	import torch

	from groot.vla.data.schema import DatasetMetadata, StateActionMetadata
	from groot.vla.data.transform.base import InvertibleModalityTransform, ModalityTransform


	class ConcatTransform(InvertibleModalityTransform):
	"""
	Concatenate the keys according to specified order.
	"""

	# -- We inherit from ModalityTransform, so we keep apply_to as well --
	apply_to: list[str] = Field(
	default_factory=list, description="Not used in this transform, kept for compatibility."
	)

	video_concat_order: list[str] = Field(
	...,
	description="Concatenation order for each video modality. "
	"Format: ['video.ego_view_pad_res224_freq20', ...]",
	)

	state_concat_order: Optional[list[str]] = Field(
	default=None,
	description="Concatenation order for each state modality. "
	"Format: ['state.position', 'state.velocity', ...].",
	)

	action_concat_order: Optional[list[str]] = Field(
	default=None,
	description="Concatenation order for each action modality. "
	"Format: ['action.position', 'action.velocity', ...].",
	)

	action_dims: dict[str, int] = Field(
	default_factory=dict,
	description="The dimensions of the action keys.",
	)
	state_dims: dict[str, int] = Field(
	default_factory=dict,
	description="The dimensions of the state keys.",
	)

	action_dims_post_transform: dict[str, int] = Field(
	default_factory=dict,
	description="The new dimensions of the action keys after transform is applied.",
	)
	state_dims_post_transform: dict[str, int] = Field(
	default_factory=dict,
	description="The new dimensions of the state keys after transform is applied.",
	)
	# Store the transform pipeline to examine for dimension changes
	_transform_pipeline: List[ModalityTransform] = PrivateAttr(default_factory=list)

	def model_dump(self, args, *kwargs):
	if kwargs.get("mode", "python") == "json":
	include = {
	"apply_to",
	"video_concat_order",
	"state_concat_order",
	"action_concat_order",
	}
	else:
	include = kwargs.pop("include", None)

	return super().model_dump(args, include=include, *kwargs)

	def set_transform_pipeline(self, transforms: List[ModalityTransform]):
	"""Set the transform pipeline so this transform can examine it for dimension changes."""
	self._transform_pipeline = transforms

	def _get_target_rotations_from_pipeline(self) -> Dict[str, str]:
	"""Extract target_rotations from StateActionTransform instances in the pipeline."""
	target_rotations = {}
	for transform in self._transform_pipeline:
	if hasattr(transform, "target_rotations"):
	transform_target_rotations = getattr(transform, "target_rotations", {})
	if transform_target_rotations:
	target_rotations.update(transform_target_rotations)
	return target_rotations

	def apply(self, data: Dict[str, Any]) -> Dict[str, Any]:
	grouped_keys = {}
	for key in data.keys():
	try:
	modality, _ = key.split(".")
	except: # noqa: E722
	### Handle language annotation special case
	if "annotation" in key:
	modality = "language"
	else:
	modality = "others"
	if modality not in grouped_keys:
	grouped_keys[modality] = []
	grouped_keys[modality].append(key)

	if "video" in grouped_keys:
	# Check if keys in video_concat_order, state_concat_order, action_concat_order are
	# ineed contained in the data. If not, then the keys are misspecified
	video_keys = grouped_keys["video"]
	assert self.video_concat_order is not None, f"{self.video_concat_order=}, {video_keys=}"
	assert all(
	item in video_keys for item in self.video_concat_order
	), f"keys in video_concat_order are misspecified, \n{video_keys=}, \n{self.video_concat_order=}"

	# Process each video view
	unsqueezed_videos = []
	for video_key in self.video_concat_order:
	video_data = data.pop(video_key)
	unsqueezed_video = np.expand_dims(
	video_data, axis=-4
	) # [..., H, W, C] -> [..., 1, H, W, C]
	unsqueezed_videos.append(unsqueezed_video)
	# Concatenate along the new axis
	unsqueezed_video = np.concatenate(unsqueezed_videos, axis=-4) # [..., V, H, W, C]

	# Video
	data["video"] = unsqueezed_video

	# "state"
	if "state" in grouped_keys:
	state_keys = grouped_keys["state"]
	assert self.state_concat_order is not None, f"{self.state_concat_order=}"
	assert all(
	item in state_keys for item in self.state_concat_order
	), f"keys in state_concat_order are misspecified, \n{state_keys=}, \n{self.state_concat_order=}"
	# Check the state dims
	for key in self.state_concat_order:
	target_shapes = [self.state_dims[key]]
	if self.is_rotation_key(key):
	target_shapes.extend(
	[3, 4, 6]
	) # 3 -> axis_angle, 4 -> quaternion, 6 -> rotation_6d
	target_shapes.append(self.state_dims[key] * 2) # Allow for sin-cos transform
	assert (
	data[key].shape[-1] in target_shapes
	), f"State dim mismatch for {key=}, {data[key].shape[-1]=}, {target_shapes=}"
	# Concatenate the state keys
	# We'll have StateActionToTensor before this transform, so here we use torch.cat
	data["state"] = torch.cat(
	[data.pop(key) for key in self.state_concat_order], dim=-1
	) # [T, D_state]

	if "action" in grouped_keys:
	action_keys = grouped_keys["action"]
	assert self.action_concat_order is not None, f"{self.action_concat_order=}"
	# Check if all keys in concat_order are present
	assert set(self.action_concat_order) == set(
	action_keys
	), f"{set(self.action_concat_order)=}, {set(action_keys)=}"
	# Record the action dims
	for key in self.action_concat_order:
	target_shapes = [self.action_dims[key]]
	if self.is_rotation_key(key):
	target_shapes.extend(
	[3, 4, 6]
	) # 3 -> axis_angle, 4 -> quaternion, 6 -> rotation_6d
	assert (
	data[key].shape[-1] in target_shapes
	), f"Action dim mismatch for {key=}, {data[key].shape[-1]=}, {target_shapes=}"
	# Concatenate the action keys
	# We'll have StateActionToTensor before this transform, so here we use torch.cat
	data["action"] = torch.cat(
	[data.pop(key) for key in self.action_concat_order], dim=-1
	) # [T, D_action]

	return data

	def unapply(self, data: dict) -> dict:
	start_dim = 0
	assert "action" in data, f"{data.keys()=}"
	# For those dataset without actions (LAPA), we'll never run unapply
	assert self.action_concat_order is not None, f"{self.action_concat_order=}"
	action_tensor = data.pop("action")
	for key in self.action_concat_order:
	if key not in self.action_dims:
	raise ValueError(f"Action dim {key} not found in action_dims.")
	end_dim = start_dim + self.get_state_action_dims_post_transform(key)
	data[key] = action_tensor[..., start_dim:end_dim]
	start_dim = end_dim
	if "state" in data:
	assert self.state_concat_order is not None, f"{self.state_concat_order=}"
	start_dim = 0
	state_tensor = data.pop("state")
	for key in self.state_concat_order:
	end_dim = start_dim + self.get_state_action_dims_post_transform(key)
	data[key] = state_tensor[..., start_dim:end_dim]
	start_dim = end_dim
	return data

	def __call__(self, data: dict) -> dict:
	return self.apply(data)

	def get_modality_metadata(self, key: str) -> StateActionMetadata:
	modality, subkey = key.split(".")
	assert self.dataset_metadata is not None, "Metadata not set"
	modality_config = getattr(self.dataset_metadata.modalities, modality)
	assert subkey in modality_config, f"{subkey=} not found in {modality_config=}"
	assert isinstance(
	modality_config[subkey], StateActionMetadata
	), f"Expected {StateActionMetadata} for {subkey=}, got {type(modality_config[subkey])=}"
	return modality_config[subkey]

	def get_state_action_dims(self, key: str) -> int:
	"""Get the dimension of a state or action key from the dataset metadata."""
	modality_config = self.get_modality_metadata(key)
	shape = modality_config.shape
	assert len(shape) == 1, f"{shape=}"
	return shape[0]

	def get_state_action_dims_post_transform(self, key: str) -> int:
	"""
	This function is used to get the dims of the state/action keys after transform is applied.
	It is different from the `get_state_action_dims` function, because this function accounts for
	the case where we apply transforms and the # of dims is change eg. after applying axis_angle transform on
	quaternion, the dims change from 4D to 3D.
	"""
	modality_config = self.get_modality_metadata(key)
	shape = modality_config.shape
	assert len(shape) == 1, f"{shape=}"

	if self.is_rotation_key(key):
	target_rotations = self._get_target_rotations_from_pipeline()
	if key in target_rotations:
	target_rotation = target_rotations[key]
	if target_rotation == "axis_angle":
	return 3
	elif target_rotation == "quaternion":
	return 4
	elif target_rotation == "rotation_6d":
	return 6
	elif target_rotation == "euler_angles":
	return 3
	else:
	raise ValueError(f"Unknown target rotation type: {target_rotation}")
	else:
	# No target rotation specified, return original dimension
	return shape[0]
	else:
	return shape[0]

	def is_rotation_key(self, key: str) -> bool:
	modality_config = self.get_modality_metadata(key)
	return modality_config.rotation_type is not None

	def set_metadata(self, dataset_metadata: DatasetMetadata):
	"""Set the metadata and compute the dimensions of the state and action keys."""
	super().set_metadata(dataset_metadata)
	# Pre-compute the dimensions of the state and action keys
	if self.action_concat_order is not None:
	for key in self.action_concat_order:
	self.action_dims[key] = self.get_state_action_dims(key)
	if self.state_concat_order is not None:
	for key in self.state_concat_order:
	self.state_dims[key] = self.get_state_action_dims(key)