from collections.abc import Sequence import logging import pathlib import time from typing import Any, TypeAlias import flax import flax.traverse_util import jax import jax.numpy as jnp import numpy as np from openpi_client import base_policy as _base_policy import torch from typing_extensions import override from openpi import transforms as _transforms from openpi.models import model as _model from openpi.shared import array_typing as at from openpi.shared import nnx_utils BasePolicy: TypeAlias = _base_policy.BasePolicy class Policy(BasePolicy): def __init__( self, model: _model.BaseModel, *, rng: at.KeyArrayLike | None = None, transforms: Sequence[_transforms.DataTransformFn] = (), output_transforms: Sequence[_transforms.DataTransformFn] = (), sample_kwargs: dict[str, Any] | None = None, metadata: dict[str, Any] | None = None, pytorch_device: str = "cpu", is_pytorch: bool = False, ): """Initialize the Policy. Args: model: The model to use for action sampling. rng: Random number generator key for JAX models. Ignored for PyTorch models. transforms: Input data transformations to apply before inference. output_transforms: Output data transformations to apply after inference. sample_kwargs: Additional keyword arguments to pass to model.sample_actions. metadata: Additional metadata to store with the policy. pytorch_device: Device to use for PyTorch models (e.g., "cpu", "cuda:0"). Only relevant when is_pytorch=True. is_pytorch: Whether the model is a PyTorch model. If False, assumes JAX model. """ self._model = model self._input_transform = _transforms.compose(transforms) self._output_transform = _transforms.compose(output_transforms) self._sample_kwargs = sample_kwargs or {} self._metadata = metadata or {} self._is_pytorch_model = is_pytorch self._pytorch_device = pytorch_device if self._is_pytorch_model: self._model = self._model.to(pytorch_device) self._model.eval() self._sample_actions = model.sample_actions else: # JAX model setup self._sample_actions = nnx_utils.module_jit(model.sample_actions) self._rng = rng or jax.random.key(0) @override def infer(self, obs: dict, *, noise: np.ndarray | None = None) -> dict: # type: ignore[misc] rtc_options = {k: obs[k] for k in obs if isinstance(k, str) and k.startswith("_rtc/")} policy_obs = {k: v for k, v in obs.items() if not (isinstance(k, str) and k.startswith("_rtc/"))} # Make a copy since transformations may modify the inputs in place. inputs = jax.tree.map(lambda x: x, policy_obs) inputs = self._input_transform(inputs) if not self._is_pytorch_model: # Make a batch and convert to jax.Array. inputs = jax.tree.map(lambda x: jnp.asarray(x)[np.newaxis, ...], inputs) self._rng, sample_rng_or_pytorch_device = jax.random.split(self._rng) else: # Convert inputs to PyTorch tensors and move to correct device inputs = jax.tree.map(lambda x: torch.from_numpy(np.array(x)).to(self._pytorch_device)[None, ...], inputs) sample_rng_or_pytorch_device = self._pytorch_device # Prepare kwargs for sample_actions sample_kwargs = dict(self._sample_kwargs) if noise is not None: noise = torch.from_numpy(noise).to(self._pytorch_device) if self._is_pytorch_model else jnp.asarray(noise) if noise.ndim == 2: # If noise is (action_horizon, action_dim), add batch dimension noise = noise[None, ...] # Make it (1, action_horizon, action_dim) sample_kwargs["noise"] = noise rtc_enabled = bool(int(rtc_options.get("_rtc/enabled", 0))) if rtc_options else False rtc_prev_actions = rtc_options.get("_rtc/previous_actions") if rtc_enabled and rtc_prev_actions is not None: if not self._is_pytorch_model: raise NotImplementedError("RTC guidance is currently implemented for PyTorch policies only.") rtc_inputs = jax.tree.map(lambda x: x, policy_obs) rtc_inputs["actions"] = np.array(rtc_prev_actions, dtype=np.float32, copy=True) rtc_inputs = self._input_transform(rtc_inputs) sample_kwargs.update( { "rtc_prev_actions": torch.from_numpy(np.array(rtc_inputs["actions"], dtype=np.float32, copy=True)).to( self._pytorch_device ), "rtc_inference_delay": int(rtc_options.get("_rtc/inference_delay", 0)), "rtc_execution_horizon": int(rtc_options.get("_rtc/execution_horizon", 0)), "rtc_max_guidance_weight": float(rtc_options.get("_rtc/max_guidance_weight", 5.0)), "rtc_schedule": str(rtc_options.get("_rtc/schedule", "exp")), "rtc_guidance_dims": int(rtc_options.get("_rtc/guidance_dims", 0)), } ) observation = _model.Observation.from_dict(inputs) pytorch_device = ( torch.device(self._pytorch_device) if self._is_pytorch_model and self._pytorch_device is not None else None ) if pytorch_device is not None and pytorch_device.type == "cuda" and torch.cuda.is_available(): torch.cuda.synchronize(pytorch_device) start_time = time.monotonic() outputs = { "state": inputs["state"], "actions": self._sample_actions(sample_rng_or_pytorch_device, observation, **sample_kwargs), } if pytorch_device is not None and pytorch_device.type == "cuda" and torch.cuda.is_available(): torch.cuda.synchronize(pytorch_device) model_time = time.monotonic() - start_time if self._is_pytorch_model: outputs = jax.tree.map(lambda x: np.asarray(x[0, ...].detach().cpu()), outputs) else: outputs = jax.tree.map(lambda x: np.asarray(x[0, ...]), outputs) outputs = self._output_transform(outputs) sample_timing = getattr(self._model, "_last_sample_timing", {}) if self._is_pytorch_model else {} outputs["policy_timing"] = { "infer_ms": model_time * 1000, **sample_timing, } return outputs @property def metadata(self) -> dict[str, Any]: return self._metadata class PolicyRecorder(_base_policy.BasePolicy): """Records the policy's behavior to disk.""" def __init__(self, policy: _base_policy.BasePolicy, record_dir: str): self._policy = policy logging.info(f"Dumping policy records to: {record_dir}") self._record_dir = pathlib.Path(record_dir) self._record_dir.mkdir(parents=True, exist_ok=True) self._record_step = 0 @override def infer(self, obs: dict) -> dict: # type: ignore[misc] results = self._policy.infer(obs) data = {"inputs": obs, "outputs": results} data = flax.traverse_util.flatten_dict(data, sep="/") output_path = self._record_dir / f"step_{self._record_step}" self._record_step += 1 np.save(output_path, np.asarray(data)) return results