Add isolated Minecraft and RE10K baseline evaluation suite

59630ba verified about 1 month ago

6.28 kB

"""
Pose and action utility functions adapted from WorldMem/framear.
"""
import numpy as np
import torch
from einops import rearrange


ACTION_KEYS = [
    "inventory",    # 0
    "ESC",          # 1
    "hotbar.1",     # 2
    "hotbar.2",     # 3
    "hotbar.3",     # 4
    "hotbar.4",     # 5
    "hotbar.5",     # 6
    "hotbar.6",     # 7
    "hotbar.7",     # 8
    "hotbar.8",     # 9
    "hotbar.9",     # 10
    "forward",      # 11
    "back",         # 12
    "left",         # 13
    "right",        # 14
    "cameraY",      # 15
    "cameraX",      # 16
    "jump",         # 17
    "sneak",        # 18
    "sprint",       # 19
    "swapHands",    # 20
    "attack",       # 21
    "use",          # 22
    "pickItem",     # 23
    "drop",         # 24
]


def convert_action_space(actions: np.ndarray) -> torch.Tensor:
    """Convert raw WorldMem action array to 25-dim one-hot vector.
    Args:
        actions: np.ndarray of shape (T, 8+), raw action indices from NPZ.
    Returns:
        Tensor of shape (T, 25) with float values.
    """
    vec_25 = torch.zeros(len(actions), len(ACTION_KEYS))
    vec_25[actions[:, 0] == 1, 11] = 1   # forward
    vec_25[actions[:, 0] == 2, 12] = 1   # backward
    vec_25[actions[:, 4] == 11, 16] = -1  # camera X left
    vec_25[actions[:, 4] == 13, 16] = 1   # camera X right
    vec_25[actions[:, 3] == 11, 15] = -1  # camera Y up
    vec_25[actions[:, 3] == 13, 15] = 1   # camera Y down
    vec_25[actions[:, 5] == 6, 24] = 1   # place box
    vec_25[actions[:, 5] == 1, 24] = 1   # place box
    vec_25[actions[:, 1] == 1, 13] = 1   # left
    vec_25[actions[:, 1] == 2, 14] = 1   # right
    vec_25[actions[:, 7] == 1, 2] = 1    # hotbar 1
    return vec_25


def euler_to_rotation_matrix(pitch: torch.Tensor, yaw: torch.Tensor) -> torch.Tensor:
    """Convert pitch and yaw angles (radians) to 3x3 rotation matrix.
    Args:
        pitch: Tensor of pitch angles, shape (...,).
        yaw:   Tensor of yaw angles, shape (...,).
    Returns:
        Rotation matrix of shape (..., 3, 3).
    """
    cos_pitch, sin_pitch = torch.cos(pitch), torch.sin(pitch)
    cos_yaw, sin_yaw = torch.cos(yaw), torch.sin(yaw)

    R_pitch = torch.stack([
        torch.ones_like(pitch), torch.zeros_like(pitch), torch.zeros_like(pitch),
        torch.zeros_like(pitch), cos_pitch, -sin_pitch,
        torch.zeros_like(pitch), sin_pitch, cos_pitch
    ], dim=-1).reshape(-1, 3, 3)

    R_yaw = torch.stack([
        cos_yaw, torch.zeros_like(yaw), sin_yaw,
        torch.zeros_like(yaw), torch.ones_like(yaw), torch.zeros_like(yaw),
        -sin_yaw, torch.zeros_like(yaw), cos_yaw
    ], dim=-1).reshape(-1, 3, 3)

    return torch.matmul(R_yaw, R_pitch)


def euler_to_camera_to_world_matrix(pose: torch.Tensor) -> torch.Tensor:
    """Convert (x, y, z, pitch, yaw) pose vector to 4x4 camera-to-world matrix.
    Args:
        pose: Tensor of shape (5,), (f, 5), or (f, b, 5).
    Returns:
        Tensor of shape (4,4), (f, 4, 4), or (f, b, 4, 4).
    """
    origin_dim = pose.ndim
    if origin_dim == 1:
        pose = pose.unsqueeze(0).unsqueeze(0)
    elif origin_dim == 2:
        pose = pose.unsqueeze(0)

    x, y, z, pitch, yaw = (
        pose[..., 0], pose[..., 1], pose[..., 2],
        pose[..., 3], pose[..., 4]
    )
    pitch = torch.deg2rad(pitch)
    yaw = torch.deg2rad(yaw)

    R = euler_to_rotation_matrix(pitch, yaw)

    eye = torch.eye(4, dtype=pose.dtype, device=pose.device)
    c2w = eye.repeat(R.shape[0], 1, 1)
    c2w[:, :3, :3] = R
    c2w[:, :3, 3] = torch.stack([x.reshape(-1), y.reshape(-1), z.reshape(-1)], dim=-1)

    if origin_dim == 3:
        return c2w.view(pose.shape[0], pose.shape[1], 4, 4)
    elif origin_dim == 2:
        return c2w.view(pose.shape[0], 4, 4)
    else:
        return c2w.squeeze(0).squeeze(0)


def c2w_to_w2c(c2w: torch.Tensor) -> torch.Tensor:
    """Invert camera-to-world to world-to-camera matrix.
    Args:
        c2w: Tensor of shape (..., 4, 4).
    Returns:
        w2c: Tensor of shape (..., 4, 4).
    """
    R = c2w[..., :3, :3]
    t = c2w[..., :3, 3]
    Rt = R.transpose(-1, -2)
    t_new = -Rt @ t.unsqueeze(-1)
    w2c = torch.zeros_like(c2w)
    w2c[..., :3, :3] = Rt
    w2c[..., :3, 3] = t_new.squeeze(-1)
    w2c[..., 3, 3] = 1.0
    return w2c


def extract_pose_action_from_npz(
    data_path: str,
    start_frame: int,
    num_frames: int,
    dtype=torch.float32
) -> tuple:
    """Load action and pose from WorldMem NPZ file.
    Args:
        data_path:   Path to .npz file containing 'actions' and 'poses'.
        start_frame: Starting frame index.
        num_frames:  Number of frames to load.
        dtype:       Output tensor dtype.
    Returns:
        actions: Tensor (num_frames, 25) - encoded action vectors.
        poses:   Tensor (num_frames, 5)  - raw pose vectors (x,y,z,pitch,yaw).
    """
    data = np.load(data_path)
    actions_raw = data["actions"]          # (N, 8+)
    poses_pool = data["poses"].copy()      # (N, 5)

    # Fix potentially corrupted height in first frame
    poses_pool[0, 1] = poses_pool[1, 1]

    # Pad poses if shorter than actions
    if len(poses_pool) < len(actions_raw):
        poses_pool = np.pad(poses_pool, ((1, 0), (0, 0)))

    actions_raw = actions_raw[start_frame: start_frame + num_frames]
    poses = poses_pool[start_frame: start_frame + num_frames]

    # Normalize pose: translate relative to first frame, negate yaw
    poses[:, :3] -= poses[:1, :3]
    poses[:, -1] = -poses[:, -1]
    poses[:, 3:] %= 360

    actions = convert_action_space(actions_raw)  # (T, 25)
    return actions.to(dtype), torch.tensor(poses, dtype=dtype)


def poses_to_viewmats(poses: torch.Tensor) -> torch.Tensor:
    """Convert (B, T, 5) pose vectors to (B, T, 4, 4) world-to-camera matrices.
    Args:
        poses: Tensor of shape (B, T, 5).
    Returns:
        viewmats: Tensor of shape (B, T, 4, 4).
    """
    B, T, _ = poses.shape
    # (B, T, 5) -> (T, B, 5) for euler_to_camera_to_world_matrix
    pose_tbd = rearrange(poses, "b t d -> t b d")
    c2w = euler_to_camera_to_world_matrix(pose_tbd)  # (T, B, 4, 4)
    w2c = c2w_to_w2c(c2w)                            # (T, B, 4, 4)
    viewmats = rearrange(w2c, "t b p q -> b t p q")  # (B, T, 4, 4)
    return viewmats

	"""
	Pose and action utility functions adapted from WorldMem/framear.
	"""
	import numpy as np
	import torch
	from einops import rearrange


	ACTION_KEYS = [
	"inventory", # 0
	"ESC", # 1
	"hotbar.1", # 2
	"hotbar.2", # 3
	"hotbar.3", # 4
	"hotbar.4", # 5
	"hotbar.5", # 6
	"hotbar.6", # 7
	"hotbar.7", # 8
	"hotbar.8", # 9
	"hotbar.9", # 10
	"forward", # 11
	"back", # 12
	"left", # 13
	"right", # 14
	"cameraY", # 15
	"cameraX", # 16
	"jump", # 17
	"sneak", # 18
	"sprint", # 19
	"swapHands", # 20
	"attack", # 21
	"use", # 22
	"pickItem", # 23
	"drop", # 24
	]


	def convert_action_space(actions: np.ndarray) -> torch.Tensor:
	"""Convert raw WorldMem action array to 25-dim one-hot vector.

	Args:
	actions: np.ndarray of shape (T, 8+), raw action indices from NPZ.

	Returns:
	Tensor of shape (T, 25) with float values.
	"""
	vec_25 = torch.zeros(len(actions), len(ACTION_KEYS))
	vec_25[actions[:, 0] == 1, 11] = 1 # forward
	vec_25[actions[:, 0] == 2, 12] = 1 # backward
	vec_25[actions[:, 4] == 11, 16] = -1 # camera X left
	vec_25[actions[:, 4] == 13, 16] = 1 # camera X right
	vec_25[actions[:, 3] == 11, 15] = -1 # camera Y up
	vec_25[actions[:, 3] == 13, 15] = 1 # camera Y down
	vec_25[actions[:, 5] == 6, 24] = 1 # place box
	vec_25[actions[:, 5] == 1, 24] = 1 # place box
	vec_25[actions[:, 1] == 1, 13] = 1 # left
	vec_25[actions[:, 1] == 2, 14] = 1 # right
	vec_25[actions[:, 7] == 1, 2] = 1 # hotbar 1
	return vec_25


	def euler_to_rotation_matrix(pitch: torch.Tensor, yaw: torch.Tensor) -> torch.Tensor:
	"""Convert pitch and yaw angles (radians) to 3x3 rotation matrix.

	Args:
	pitch: Tensor of pitch angles, shape (...,).
	yaw: Tensor of yaw angles, shape (...,).

	Returns:
	Rotation matrix of shape (..., 3, 3).
	"""
	cos_pitch, sin_pitch = torch.cos(pitch), torch.sin(pitch)
	cos_yaw, sin_yaw = torch.cos(yaw), torch.sin(yaw)

	R_pitch = torch.stack([
	torch.ones_like(pitch), torch.zeros_like(pitch), torch.zeros_like(pitch),
	torch.zeros_like(pitch), cos_pitch, -sin_pitch,
	torch.zeros_like(pitch), sin_pitch, cos_pitch
	], dim=-1).reshape(-1, 3, 3)

	R_yaw = torch.stack([
	cos_yaw, torch.zeros_like(yaw), sin_yaw,
	torch.zeros_like(yaw), torch.ones_like(yaw), torch.zeros_like(yaw),
	-sin_yaw, torch.zeros_like(yaw), cos_yaw
	], dim=-1).reshape(-1, 3, 3)

	return torch.matmul(R_yaw, R_pitch)


	def euler_to_camera_to_world_matrix(pose: torch.Tensor) -> torch.Tensor:
	"""Convert (x, y, z, pitch, yaw) pose vector to 4x4 camera-to-world matrix.

	Args:
	pose: Tensor of shape (5,), (f, 5), or (f, b, 5).

	Returns:
	Tensor of shape (4,4), (f, 4, 4), or (f, b, 4, 4).
	"""
	origin_dim = pose.ndim
	if origin_dim == 1:
	pose = pose.unsqueeze(0).unsqueeze(0)
	elif origin_dim == 2:
	pose = pose.unsqueeze(0)

	x, y, z, pitch, yaw = (
	pose[..., 0], pose[..., 1], pose[..., 2],
	pose[..., 3], pose[..., 4]
	)
	pitch = torch.deg2rad(pitch)
	yaw = torch.deg2rad(yaw)

	R = euler_to_rotation_matrix(pitch, yaw)

	eye = torch.eye(4, dtype=pose.dtype, device=pose.device)
	c2w = eye.repeat(R.shape[0], 1, 1)
	c2w[:, :3, :3] = R
	c2w[:, :3, 3] = torch.stack([x.reshape(-1), y.reshape(-1), z.reshape(-1)], dim=-1)

	if origin_dim == 3:
	return c2w.view(pose.shape[0], pose.shape[1], 4, 4)
	elif origin_dim == 2:
	return c2w.view(pose.shape[0], 4, 4)
	else:
	return c2w.squeeze(0).squeeze(0)


	def c2w_to_w2c(c2w: torch.Tensor) -> torch.Tensor:
	"""Invert camera-to-world to world-to-camera matrix.

	Args:
	c2w: Tensor of shape (..., 4, 4).

	Returns:
	w2c: Tensor of shape (..., 4, 4).
	"""
	R = c2w[..., :3, :3]
	t = c2w[..., :3, 3]
	Rt = R.transpose(-1, -2)
	t_new = -Rt @ t.unsqueeze(-1)
	w2c = torch.zeros_like(c2w)
	w2c[..., :3, :3] = Rt
	w2c[..., :3, 3] = t_new.squeeze(-1)
	w2c[..., 3, 3] = 1.0
	return w2c


	def extract_pose_action_from_npz(
	data_path: str,
	start_frame: int,
	num_frames: int,
	dtype=torch.float32
	) -> tuple:
	"""Load action and pose from WorldMem NPZ file.

	Args:
	data_path: Path to .npz file containing 'actions' and 'poses'.
	start_frame: Starting frame index.
	num_frames: Number of frames to load.
	dtype: Output tensor dtype.

	Returns:
	actions: Tensor (num_frames, 25) - encoded action vectors.
	poses: Tensor (num_frames, 5) - raw pose vectors (x,y,z,pitch,yaw).
	"""
	data = np.load(data_path)
	actions_raw = data["actions"] # (N, 8+)
	poses_pool = data["poses"].copy() # (N, 5)

	# Fix potentially corrupted height in first frame
	poses_pool[0, 1] = poses_pool[1, 1]

	# Pad poses if shorter than actions
	if len(poses_pool) < len(actions_raw):
	poses_pool = np.pad(poses_pool, ((1, 0), (0, 0)))

	actions_raw = actions_raw[start_frame: start_frame + num_frames]
	poses = poses_pool[start_frame: start_frame + num_frames]

	# Normalize pose: translate relative to first frame, negate yaw
	poses[:, :3] -= poses[:1, :3]
	poses[:, -1] = -poses[:, -1]
	poses[:, 3:] %= 360

	actions = convert_action_space(actions_raw) # (T, 25)
	return actions.to(dtype), torch.tensor(poses, dtype=dtype)


	def poses_to_viewmats(poses: torch.Tensor) -> torch.Tensor:
	"""Convert (B, T, 5) pose vectors to (B, T, 4, 4) world-to-camera matrices.

	Args:
	poses: Tensor of shape (B, T, 5).

	Returns:
	viewmats: Tensor of shape (B, T, 4, 4).
	"""
	B, T, _ = poses.shape
	# (B, T, 5) -> (T, B, 5) for euler_to_camera_to_world_matrix
	pose_tbd = rearrange(poses, "b t d -> t b d")
	c2w = euler_to_camera_to_world_matrix(pose_tbd) # (T, B, 4, 4)
	w2c = c2w_to_w2c(c2w) # (T, B, 4, 4)
	viewmats = rearrange(w2c, "t b p q -> b t p q") # (B, T, 4, 4)
	return viewmats