dovla_cil/models/action_encoder.py

from __future__ import annotations

from typing import Any

from dovla_cil.data.schema import ActionChunk

try:
    import torch
    from torch import nn
except ImportError:  # pragma: no cover - torch is an install dependency, absent in bare smoke envs
    torch = None
    nn = None


TOY_COMMANDS = (
    "noop",
    "move_to",
    "grasp",
    "release",
    "push",
    "place_at",
    "open",
    "close",
    "delay",
)
TOY_SKILLS = (
    "unknown",
    "noop",
    "delay",
    "reach",
    "grasp",
    "lift",
    "push",
    "place",
    "place_at",
    "open",
    "close",
    "opened",
    "closed",
)


def vectorize_toy_action(
    action: ActionChunk | dict[str, Any] | list[dict[str, Any]] | list[list[float]],
    *,
    action_dim: int = 8,
    action_horizon: int = 4,
) -> list[list[float]]:
    """Convert an action chunk into a fixed horizon numeric matrix.

    Symbolic toy actions use intentionally simple and stable feature slots:
    command id, target hash, reference hash, dx, dy, x, y, z, then any extra slots remain zero.
    Numeric simulator action chunks, such as ManiSkill control vectors, are preserved directly and
    only padded or truncated to the requested shape.
    """

    if action_dim <= 0:
        raise ValueError("action_dim must be positive")
    if action_horizon <= 0:
        raise ValueError("action_horizon must be positive")
    numeric_rows = _numeric_action_rows(action)
    if numeric_rows is not None:
        return _pad_numeric_rows(numeric_rows, action_dim=action_dim, action_horizon=action_horizon)
    commands = _action_commands(action)
    rows: list[list[float]] = []
    for command in commands[:action_horizon]:
        row = [0.0] * action_dim
        if action_dim > 0:
            row[0] = _command_id(str(command.get("command") or command.get("type") or "noop"))
        if action_dim > 1:
            row[1] = _stable_unit_hash(str(command.get("object") or command.get("target") or ""))
        if action_dim > 2:
            row[2] = _stable_unit_hash(
                str(command.get("reference") or command.get("container") or "")
            )
        if action_dim > 3:
            row[3] = float(command.get("dx", 0.0) or 0.0)
        if action_dim > 4:
            row[4] = float(command.get("dy", 0.0) or 0.0)
        if action_dim > 5:
            position = _position(command.get("position", [0.0, 0.0, 0.0]))
            row[5] = position[0]
        if action_dim > 6:
            row[6] = position[1]
        if action_dim > 7:
            row[7] = position[2]
        rows.append(row)
    while len(rows) < action_horizon:
        rows.append([0.0] * action_dim)
    return rows


def devectorize_toy_action(
    values: list[float] | list[list[float]],
    *,
    skill_type: str | None = None,
) -> ActionChunk:
    """Best-effort conversion from a numeric policy output to a toy ActionChunk."""

    rows = _matrix(values)
    commands: list[dict[str, Any]] = []
    for row in rows:
        if not row or all(abs(float(value)) < 1e-8 for value in row):
            continue
        command = _nearest_command(float(row[0]))
        if command in {"noop", "delay"}:
            payload: dict[str, Any] = {"command": command}
            if command == "delay":
                payload["steps"] = max(1, int(round(abs(row[1]) * 3)) if len(row) > 1 else 1)
        elif command == "push":
            payload = {
                "command": "push",
                "object": "predicted_target",
                "dx": float(row[3]) if len(row) > 3 else 0.0,
                "dy": float(row[4]) if len(row) > 4 else 0.0,
            }
        elif command in {"move_to", "place_at"}:
            payload = {
                "command": command,
                "object": "predicted_target",
                "position": [
                    float(row[5]) if len(row) > 5 else 0.0,
                    float(row[6]) if len(row) > 6 else 0.0,
                    float(row[7]) if len(row) > 7 else 0.03,
                ],
            }
        elif command in {"grasp", "open", "close"}:
            payload = {"command": command, "object": "predicted_target"}
        else:
            payload = {"command": command}
        commands.append(payload)
    if not commands:
        commands = [{"command": "noop"}]
    return ActionChunk(
        representation="semantic",
        horizon=len(commands),
        values=commands,
        skill_type=skill_type,
        metadata={"source": "devectorized_toy_policy"},
    )


def skill_type_id(skill_type: str | None) -> int:
    normalized = str(skill_type or "unknown")
    try:
        return TOY_SKILLS.index(normalized)
    except ValueError:
        return 0


if nn is not None:

    class ActionEncoder(nn.Module):
        def __init__(
            self,
            action_dim: int,
            hidden_dim: int,
            *,
            action_horizon: int = 4,
            skill_vocab_size: int = 32,
        ) -> None:
            super().__init__()
            self.action_dim = int(action_dim)
            self.action_horizon = int(action_horizon)
            self.step_mlp = nn.Sequential(
                nn.Linear(self.action_dim, hidden_dim),
                nn.LayerNorm(hidden_dim),
                nn.GELU(),
                nn.Linear(hidden_dim, hidden_dim),
            )
            self.skill_embedding = nn.Embedding(skill_vocab_size, hidden_dim)
            self.out = nn.Sequential(
                nn.LayerNorm(hidden_dim),
                nn.GELU(),
                nn.Linear(hidden_dim, hidden_dim),
            )

        def forward(self, action, skill_type=None):
            action_tensor = coerce_action_tensor(
                action,
                action_dim=self.action_dim,
                action_horizon=self.action_horizon,
                device=next(self.parameters()).device,
            )
            step_features = self.step_mlp(action_tensor)
            pooled = step_features.mean(dim=1)
            if skill_type is not None:
                pooled = pooled + self.skill_embedding(
                    coerce_skill_ids(skill_type, device=pooled.device)
                    % self.skill_embedding.num_embeddings
                )
            return self.out(pooled)

else:

    class ActionEncoder:  # pragma: no cover
        def __init__(self, *args, **kwargs) -> None:
            del args, kwargs
            raise ImportError("Install torch to use ActionEncoder.")


def coerce_action_tensor(
    action,
    *,
    action_dim: int,
    action_horizon: int,
    device=None,
):
    if torch is None:  # pragma: no cover
        raise ImportError("Install torch to coerce action tensors.")
    if isinstance(action, ActionChunk):
        matrix = vectorize_toy_action(action, action_dim=action_dim, action_horizon=action_horizon)
        return torch.tensor([matrix], dtype=torch.float32, device=device)
    if isinstance(action, list) and action and isinstance(action[0], ActionChunk):
        matrices = [
            vectorize_toy_action(item, action_dim=action_dim, action_horizon=action_horizon)
            for item in action
        ]
        return torch.tensor(matrices, dtype=torch.float32, device=device)
    tensor = torch.as_tensor(action, dtype=torch.float32, device=device)
    if tensor.ndim == 1:
        tensor = tensor.unsqueeze(0)
    if tensor.ndim == 2:
        if tensor.shape[-1] == action_horizon * action_dim:
            tensor = tensor.reshape(tensor.shape[0], action_horizon, action_dim)
        else:
            tensor = tensor.unsqueeze(1)
    if tensor.ndim != 3:
        raise ValueError("Action tensor must have shape [B,D], [B,H,D], or [D]")
    return _pad_action_tensor(tensor, action_horizon=action_horizon, action_dim=action_dim)


def coerce_skill_ids(skill_type, *, device=None):
    if torch is None:  # pragma: no cover
        raise ImportError("Install torch to coerce skill ids.")
    if isinstance(skill_type, torch.Tensor):
        return skill_type.to(device=device, dtype=torch.long).flatten()
    if isinstance(skill_type, str) or skill_type is None:
        ids = [skill_type_id(skill_type)]
    else:
        ids = [skill_type_id(item) for item in skill_type]
    return torch.tensor(ids, dtype=torch.long, device=device)


def _pad_action_tensor(tensor, *, action_horizon: int, action_dim: int):
    if tensor.shape[1] > action_horizon:
        tensor = tensor[:, :action_horizon, :]
    elif tensor.shape[1] < action_horizon:
        pad = tensor.new_zeros((tensor.shape[0], action_horizon - tensor.shape[1], tensor.shape[2]))
        tensor = torch.cat([tensor, pad], dim=1)
    if tensor.shape[2] > action_dim:
        tensor = tensor[:, :, :action_dim]
    elif tensor.shape[2] < action_dim:
        pad = tensor.new_zeros((tensor.shape[0], tensor.shape[1], action_dim - tensor.shape[2]))
        tensor = torch.cat([tensor, pad], dim=2)
    return tensor


def _action_commands(
    action: ActionChunk | dict[str, Any] | list[dict[str, Any]],
) -> list[dict[str, Any]]:
    if isinstance(action, ActionChunk):
        if isinstance(action.values, list) and all(
            isinstance(item, dict) for item in action.values
        ):
            return [dict(item) for item in action.values]
        return [{"command": "numeric", "position": action.flat_values[:3]}]
    if isinstance(action, dict):
        return [dict(action)]
    return [dict(item) for item in action]


def _position(value: Any) -> list[float]:
    if isinstance(value, dict):
        return [
            float(value.get("x", 0.0)),
            float(value.get("y", 0.0)),
            float(value.get("z", 0.0)),
        ]
    if isinstance(value, list | tuple):
        items = list(value)
        while len(items) < 3:
            items.append(0.0)
        return [float(items[0]), float(items[1]), float(items[2])]
    return [0.0, 0.0, 0.0]


def _matrix(values: list[float] | list[list[float]]) -> list[list[float]]:
    if not values:
        return []
    if all(isinstance(item, int | float) for item in values):
        return [[float(item) for item in values]]  # type: ignore[arg-type]
    return [[float(item) for item in row] for row in values]  # type: ignore[union-attr]


def _numeric_action_rows(action: Any) -> list[list[float]] | None:
    values = action.values if isinstance(action, ActionChunk) else action
    if not isinstance(values, list) or not values:
        return None
    if all(isinstance(row, list) and _is_numeric_sequence(row) for row in values):
        return [[float(item) for item in row] for row in values]
    if _is_numeric_sequence(values):
        return [[float(item) for item in values]]
    return None


def _is_numeric_sequence(values: Any) -> bool:
    return isinstance(values, list) and all(isinstance(item, int | float) for item in values)


def _pad_numeric_rows(
    rows: list[list[float]], *, action_dim: int, action_horizon: int
) -> list[list[float]]:
    output: list[list[float]] = []
    for row in rows[:action_horizon]:
        clipped = [float(value) for value in row[:action_dim]]
        if len(clipped) < action_dim:
            clipped.extend([0.0] * (action_dim - len(clipped)))
        output.append(clipped)
    while len(output) < action_horizon:
        output.append([0.0] * action_dim)
    return output


def _command_id(command: str) -> float:
    try:
        return TOY_COMMANDS.index(command) / max(len(TOY_COMMANDS) - 1, 1)
    except ValueError:
        return 0.0


def _nearest_command(value: float) -> str:
    index = round(max(0.0, min(1.0, value)) * max(len(TOY_COMMANDS) - 1, 1))
    return TOY_COMMANDS[int(index)]


def _stable_unit_hash(value: str) -> float:
    total = 0
    for byte in value.encode("utf-8"):
        total = (total * 257 + byte) % 1000003
    return total / 1000003.0