sim/simulator.py

import cv2
import torch
import numpy as np
import einops
import skimage
import time
import robomimic.utils.file_utils as FileUtils

from sim.robomimic.robomimic_wrapper import RobomimicLowdimWrapper
from sim.robomimic.robomimic_runner import create_env, OBS_KEYS, RESOLUTION
from genie.st_mask_git import STMaskGIT
from genie.st_mar import STMAR
from datasets.utils import get_image_encoder
from data import DATA_FREQ_TABLE
from train_diffusion import SVD_SCALE

from typing import Optional, Tuple, Callable, Dict


class Simulator:
    def set_initial_state(self, state):
        """
        the initial state of the simulated scene
        e.g.
        1. in robomimic, it's the scene state vector
        2. in genie, it's the initial frames to prompt the model
        """
        raise NotImplementedError

    @torch.inference_mode()
    def step(self, action):
        raise NotImplementedError
    
    def reset(self):
        raise NotImplementedError
    
    def close(self):
        raise NotImplementedError
    
    @property
    def dt(self):
        raise NotImplementedError


class PhysicsSimulator(Simulator):
    def __init__(self):
        super().__init__()

    # physics engine should be able to update dt
    def set_dt(self, dt):
        raise NotImplementedError
    
    # physics engine should be able to get scene state
    # e.g., robot joint positions, object positions, etc.
    def get_raw_state(self, port: Optional[str] = None):
        raise NotImplementedError
    
    @property
    def action_dimension(self):
        raise NotImplementedError
    

class LearnedSimulator(Simulator):
    def __init__(self):
        super().__init__()


# data replayed respect physics, so we inherit from PhysicsSimulator
# it can be considered as a special case of PhysicsSimulator
class ReplaySimulator(PhysicsSimulator):
    def __init__(self, 
        frames, 
        prompt_horizon: int = 0, 
        dt: Optional[float] = None
    ):
        super().__init__()
        self.frames = frames
        self.frame_idx = prompt_horizon
        assert self.frame_idx < len(self.frames)
        self._dt = dt
        self.prompt_horizon = prompt_horizon
    
    def __len__(self):
        return len(self.frames) - self.prompt_horizon
    
    def step(self, action):
        frame = self.frames[self.frame_idx]
        assert self.frame_idx < len(self.frames)
        self.frame_idx = self.frame_idx + 1
        return {
            'pred_next_frame': frame
        }
    
    def reset(self):    # return current frame = last frame of prompt
        self.frame_idx = self.prompt_horizon
        return self.prompt()[-1]
    
    def prompt(self):
        return self.frames[:self.prompt_horizon]
    
    @property
    def dt(self):
        return self._dt
    

class RobomimicSimulator(PhysicsSimulator):
    """
    Example usage:
    sim = RobomimicSimulator("lift")
    useless_low_dim_state = sim.reset()
    image = sim.step([0] * 7)
    Image.fromarray(image).save("test.png")
    """
    def __init__(self, env_name):
        super().__init__()
        env_meta = FileUtils.get_env_metadata_from_dataset(
            f'data/robomimic_datasets/robomimic_raw/datasets/{env_name}/ph/image.hdf5')
        env_meta['use_image_obs'] = True
        env = create_env(env_meta, obs_keys=OBS_KEYS)
        self.env = RobomimicLowdimWrapper(env=env, obs_keys=OBS_KEYS)
        self.env_name = env_name

    def set_initial_state(self, state):
        self.env.init_state = state

    def step(self, action):
        obs, reward, done, info = self.env.step(action)
        img = self.env.render()
        img = cv2.resize(img, RESOLUTION)
        return {
            'pred_next_frame': img,
            'reward': reward,
            'done': done,
            'info': info
        }

    def reset(self):
        self.env.reset()    # reset to initial state
        img = self.env.render()
        img = cv2.resize(img, RESOLUTION)
        return img

    def close(self):
        self.env.close()

    @property
    def dt(self):
        return self.env.env.env.sim.model.opt.timestep
    
    @property
    def action_dimension(self):
        return self.env.env.action_dimension

    def set_dt(self, dt):
        self.env.env.env.sim.model.opt.timestep = dt

    def get_raw_state(self, port: Optional[str] = None):
        raw_state = self.env.env.get_state()
        raw_obs = self.env.env.get_observation()
        raw_state.update(raw_obs)
        if port is not None:
            return raw_state[port]
        return raw_state


class GenieSimulator(LearnedSimulator):

    average_delta_psnr_over = 5

    def __init__(self,
        # image preprocessing
        max_image_resolution: int = 1024,
        resize_image: bool = True,
        resize_image_resolution: int = 256,
        # tokenizer setting
        image_encoder_type: str = "temporalvae",
        image_encoder_ckpt: str = "stabilityai/stable-video-diffusion-img2vid",
        quantize: bool = False,
        quantization_slice_size: int = 16,
        # dynamics backbone setting
        backbone_type: str = "stmar",
        backbone_ckpt: str = "data/mar_ckpt/robomimic",
        prompt_horizon: int = 11,
        inference_iterations: Optional[int] = None,
        sampling_temperature: float = 0.0,
        action_stride: Optional[int] = None,
        domain: str = "robomimic",
        genie_frequency: int = 2,
        # misc
        measure_step_time: bool = False,
        compute_psnr: bool = False,
        compute_delta_psnr: bool = False,   # act as a signal for controlability
        gaussian_action_perturbation_scale: Optional[float] = None,
        device: str = 'cuda',
        physics_simulator: Optional[PhysicsSimulator] = None,
        physics_simulator_teacher_force: Optional[int] = None,
        post_processor: Optional[Callable] = None, # on the predicted image, e.g., add action
        allow_external_prompt: bool = False
    ):
        super().__init__()

        assert quantize == (image_encoder_type == "magvit"), \
            "Currently quantization if and only if magvit is the image encoder."
        assert image_encoder_type in ["magvit", "temporalvae"], \
            "Image encoder type must be either 'magvit' or 'temporalvae'."
        assert not quantize or image_encoder_type == "magvit", \
            "If quantize is enabled, image encoder type must be 'magvit'."
        assert backbone_type in ["stmaskgit", "stmar"], \
            "Backbone type must be either 'stmaskgit' or 'stmar'."
        if physics_simulator is None:
            assert physics_simulator_teacher_force is None, \
                "Physics simulator teacher force is only available when physics simulator is provided."
            assert compute_psnr is False, \
                "PSNR computation is only available when physics simulator is provided."
            assert compute_delta_psnr is False, \
                "Delta PSNR computation is only available when physics simulator is provided."

        if action_stride is None:
            action_stride = DATA_FREQ_TABLE[domain] // genie_frequency
        if compute_delta_psnr:
            compute_psnr = True # to compute delta psnr, psnr must be computed
        if inference_iterations is None:
            if backbone_type == "stmaskgit":
                inference_iterations = 2
            elif backbone_type == "stmar":
                inference_iterations = 2

        # misc     
        self.device = torch.device(device)
        self.measure_step_time = measure_step_time
        self.compute_psnr = compute_psnr
        self.compute_delta_psnr = compute_delta_psnr
        self.allow_external_prompt = allow_external_prompt

        # image preprocessing
        self.max_image_resolution = max_image_resolution
        self.resize_image = resize_image
        self.resize_image_resolution = resize_image_resolution

        # load image encoder
        self.image_encoding_dtype = torch.bfloat16
        self.quantize = quantize
        self.quant_slice_size = quantization_slice_size
        self.image_encoder_type = image_encoder_type
        self.image_encoder = get_image_encoder(
            image_encoder_type,
            image_encoder_ckpt
            ).to(device=self.device, dtype=self.image_encoding_dtype).eval()

        # load STMaskGIT model (STMAR is inherited from STMaskGIT)
        self.prompt_horizon = prompt_horizon
        self.domain = domain
        self.genie_frequency = genie_frequency
        self.inference_iterations = inference_iterations
        self.sampling_temperature = sampling_temperature
        self.action_stride = action_stride
        self.gauss_act_perturb_scale = gaussian_action_perturbation_scale
        self.backbone_type = backbone_type
        if backbone_type == "stmaskgit":
            self.backbone = STMaskGIT.from_pretrained(backbone_ckpt)
        else:
            self.backbone = STMAR.from_pretrained(backbone_ckpt)
        self.backbone = self.backbone.to(device=self.device).eval()
        
        self.post_processor = post_processor
        
        # load physics simulator if available
        # the phys sim to get ground truth image, 
        # assume the phys sim has aligned prompt frames
        self.gt_phys_sim = physics_simulator
        self.gt_teacher_force = physics_simulator_teacher_force

        # history buffer, i.e., the input to the model
        self.cached_actions = None          # (prompt_horizon, action_stride, A)
        self.cached_latent_frames = None    # (prompt_horizon, ...)
        self.init_prompt = None             # (prompt_frames, prompt_actions)

        self.step_count = 0

        # report model size
        print(
            "================ Model Size Report ================\n"
            f"    encoder size: {sum(p.numel() for p in self.image_encoder.parameters()) / 1e6:.3f}M \n"
            f"    backbone size: {sum(p.numel() for p in self.backbone.parameters()) / 1e6:.3f}M\n"
            "==================================================="
        )


    def set_initial_state(self, state: Tuple[np.ndarray, np.ndarray]):
        if not self.allow_external_prompt and self.gt_phys_sim is not None:
            raise NotImplementedError("Initial state is set by the physics simulator.")
        self.init_prompt = state


    @torch.inference_mode()
    def step(self, action: np.ndarray) -> Dict:
        # action: (action_stride, A) OR (A,)
        # return: (H, W, 3)
        assert self.cached_latent_frames is not None and self.cached_actions is not None, \
            "Model is not prompted yet. Please call `set_initial_state` first."
        
        if action.ndim == 1:
            action = np.tile(action, (self.action_stride, 1))

        # perturb action
        if self.gauss_act_perturb_scale is not None:
            action = np.random.normal(action, self.gauss_act_perturb_scale)

        # encoding
        input_latent_states = torch.cat([
            self.cached_latent_frames,
            torch.zeros_like(self.cached_latent_frames[-1:]),
        ]).unsqueeze(0).to(torch.float32)

        # dtype conversion and mask token
        if self.backbone_type == "stmaskgit":
            input_latent_states = input_latent_states.long()
            input_latent_states[:, self.prompt_horizon] = self.backbone.mask_token_id
        elif self.backbone_type == "stmar":
            input_latent_states[:, self.prompt_horizon] = self.backbone.mask_token

        # dynamics rollout
        action = torch.from_numpy(action).to(device=self.device)
        input_actions = torch.cat([     # (1, prompt_horizon + 1, action_stride * A)
            self.cached_actions,
            action.unsqueeze(0)
        ]).view(1, self.prompt_horizon + 1, -1).to(torch.float32)

        if self.measure_step_time:
            start_time = time.time()
        pred_next_latent_state = self.backbone.maskgit_generate(
            input_latent_states,
            out_t=self.prompt_horizon,
            maskgit_steps=self.inference_iterations,
            temperature=self.sampling_temperature,
            action_ids=input_actions,
            domain=[self.domain]
        )[0].squeeze(0)

        # decoding
        pred_next_frame = self._decode_image(pred_next_latent_state)

        # timing
        if self.measure_step_time:
            end_time = time.time()

        step_result = {'pred_next_frame': pred_next_frame,}
        if self.measure_step_time:
            step_result['step_time'] = end_time - start_time

        # physics simulation
        if self.gt_phys_sim is not None:
            for a in action.cpu().numpy():
                gt_result = self.gt_phys_sim.step(a)
            gt_next_frame = cv2.resize(gt_result['pred_next_frame'], pred_next_frame.shape[:2])
            step_result['gt_next_frame'] = gt_next_frame
            gt_result.pop('pred_next_frame')
            step_result.update(gt_result)

            # gt state observation
            try:
                raw_state = self.gt_phys_sim.get_raw_state()
                step_result.update(raw_state)
            except NotImplementedError:
                pass

            # compute PSNR against ground truth
            if self.compute_psnr:
                psnr = skimage.metrics.peak_signal_noise_ratio(
                    image_true=gt_next_frame / 255., 
                    image_test=pred_next_frame / 255.,
                    data_range=1.0
                )
                step_result['psnr'] = psnr

            # controlability metric
            if self.compute_delta_psnr:
                delta_psnr = 0.0
                for _ in range(self.average_delta_psnr_over):
                    # re-mask the input latent states for masked prediction
                    if self.backbone_type == "stmaskgit":
                        input_latent_states = input_latent_states.long()
                        input_latent_states[:, self.prompt_horizon] = self.backbone.mask_token_id
                    elif self.backbone_type == "stmar":
                        input_latent_states[:, self.prompt_horizon] = self.backbone.mask_token
                    # sample random action from N(0, 1)
                    random_input_actions = torch.randn_like(input_actions)
                    random_pred_next_latent_state = self.backbone.maskgit_generate(
                        input_latent_states,
                        out_t=self.prompt_horizon,
                        maskgit_steps=self.inference_iterations,
                        temperature=self.sampling_temperature,
                        action_ids=random_input_actions,
                        domain=[self.domain],
                        skip_normalization=True
                    )[0].squeeze(0)
                    random_pred_next_frame = self._decode_image(random_pred_next_latent_state)
                    this_delta_psnr = step_result['psnr'] - skimage.metrics.peak_signal_noise_ratio(
                        image_true=gt_next_frame / 255.,
                        image_test=random_pred_next_frame / 255.,
                        data_range=1.0
                    )
                    delta_psnr += this_delta_psnr / self.average_delta_psnr_over
                step_result['delta_psnr'] = delta_psnr

            if self.gt_teacher_force is not None and self.step_count % self.gt_teacher_force == 0:
                pred_next_latent_state = self._encode_image(gt_next_frame)
        
        # update history buffer
        self.cached_latent_frames = torch.cat([
            self.cached_latent_frames[1:], pred_next_latent_state.unsqueeze(0)
        ])
        self.cached_actions = torch.cat([
            self.cached_actions[1:], action.unsqueeze(0)
        ])
        
        # post processing
        if self.post_processor is not None:
            pred_next_frame = self.post_processor(pred_next_frame, action)

        self.step_count += 1

        return step_result
    

    @torch.inference_mode()
    def _encode_image(self, image: np.ndarray) -> torch.Tensor:
        # (H, W, 3)
        image = torch.from_numpy(
            self._normalize_image(image).transpose(2, 0, 1)
            ).to(device=self.device, dtype=self.image_encoding_dtype
            ).unsqueeze(0)
        H, W = image.shape[-2:]

        if self.quantize:
            H //= self.quant_slice_size
            W //= self.quant_slice_size
            _, _, indices, _ = self.image_encoder.encode(image, flip=True)
            indices = einops.rearrange(indices, "(h w) -> h w", h=H, w=W)
            indices = indices.to(torch.int32)
            return indices

        else:
            if self.image_encoder_type == "magvit":
                latent = self.image_encoder.encode_without_quantize(image)
            elif self.image_encoder_type == "temporalvae":
                latent_dist = self.image_encoder.encode(image).latent_dist
                latent = latent_dist.mean
                latent *= SVD_SCALE
                latent = einops.rearrange(latent, "b c h w -> b h w c")
            else:
                pass
            latent = latent.squeeze(0).to(torch.float32)
            return latent


    @torch.inference_mode()
    def _decode_image(self, latent: torch.Tensor) -> np.ndarray:
        # latent can be either quantized indices or raw latent
        # return (H, W, 3)

        latent = latent.to(device=self.device).unsqueeze(0)

        if self.quantize:
            latent = self.image_encoder.quantize.get_codebook_entry(
                einops.rearrange(latent, "b h w -> b (h w)"),
                bhwc=(*latent.shape, self.image_encoder.quantize.codebook_dim)
            ).flip(1)

        latent = latent.to(device=self.device, dtype=self.image_encoding_dtype)
        if self.image_encoder_type == "magvit":
            decoded_image = self.image_encoder.decode(latent)
        elif self.image_encoder_type == "temporalvae":
            latent = einops.rearrange(latent, "b h w c -> b c h w")
            latent /= SVD_SCALE
            # HACK: clip for less visual artifacts
            latent = torch.clamp(latent, -25, 25)
            decoded_image = self.image_encoder.decode(latent, num_frames=1).sample
        decoded_image = decoded_image.squeeze(0).to(torch.float32).detach().cpu().numpy()
        decoded_image = self._unnormalize_image(decoded_image).transpose(1, 2, 0)
        return decoded_image
    

    def _normalize_image(self, image: np.ndarray) -> np.ndarray:
        # (H, W, 3) normalized to [-1, 1]
        # if `resize`, resize the shorter side to `resized_res` 
        #   and then do a center crop

        image = np.asarray(image, dtype=np.float32)
        image /= 255.
        H, W = image.shape[:2]

        # resize if asked
        if self.resize_image:
            resized_res = self.resize_image_resolution  
            if H < W:
                Hnew, Wnew = resized_res, int(resized_res * W / H)
            else:
                Hnew, Wnew = int(resized_res * H / W), resized_res
            image = cv2.resize(image, (Wnew, Hnew))

            # center crop
            H, W = image.shape[:2]
            Hstart = (H - resized_res) // 2
            Wstart = (W - resized_res) // 2
            image = image[Hstart:Hstart + resized_res, Wstart:Wstart + resized_res]

        # resize if resolution is too large
        elif H > self.max_image_resolution or W > self.max_image_resolution:
            if H < W:
                Hnew, Wnew = int(self.max_image_resolution * H / W), self.max_image_resolution
            else:
                Hnew, Wnew = self.max_image_resolution, int(self.max_image_resolution * W / H)
            image = cv2.resize(image, (Wnew, Hnew))

        image = (image * 2 - 1.)
        return image


    def _unnormalize_image(self, image: np.ndarray) -> np.ndarray:
        # (H, W, 3) from [-1, 1] to [0, 255]
        # NOTE: clip happens here
        image = (image + 1.) * 127.5
        image = np.clip(image, 0, 255).astype(np.uint8)
        return image


    def reset(self) -> np.ndarray:
        # if ground truth physics simulator is provided, 
        # return the the side-by-side concatenated image

        # get the initial prompt from the physics simulator if not yet set
        if not self.allow_external_prompt and self.gt_phys_sim is not None:
            image_prompt = np.tile(
                self.gt_phys_sim.reset(), (self.prompt_horizon, 1, 1, 1)
            ).astype(np.uint8)
            action_prompt = np.zeros(
                (self.prompt_horizon, self.action_stride, self.gt_phys_sim.action_dimension)
            ).astype(np.float32)
        else: 
            assert self.init_prompt is not None, "Initial state is not set."
            image_prompt, action_prompt = self.init_prompt

        # standardize the image
        image_prompt = [self._unnormalize_image(self._normalize_image(frame)) for frame in image_prompt]

        current_image = image_prompt[-1]

        action_prompt = torch.from_numpy(action_prompt).to(device=self.device)
        self.cached_actions = action_prompt

        # convert to latent
        self.cached_latent_frames = torch.stack([
            self._encode_image(frame) for frame in image_prompt
        ], axis=0)

        if self.resize_image:
            current_image = cv2.resize(current_image, 
                (self.resize_image_resolution, self.resize_image_resolution))

        if self.gt_phys_sim is not None:
            current_image = np.concatenate([current_image, current_image], axis=1)

        self.step_count = 0

        return current_image


    def close(self):
        if self.gt_phys_sim is not None:
            try:
                self.gt_phys_sim.close()
            except NotImplementedError:
                pass


    @property
    def dt(self):
        return 1.0 / self.genie_frequency