File size: 7,628 Bytes
9a4a86b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 | from collections.abc import Sequence
import inspect
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
self._supports_prefix_z = False
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._supports_prefix_z = "return_prefix_z" in inspect.signature(model.sample_actions).parameters
if self._supports_prefix_z:
self._sample_actions = nnx_utils.module_jit(
model.sample_actions,
static_argnames=("return_prefix_z",),
)
else:
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]
# Make a copy since transformations may modify the inputs in place.
inputs = jax.tree.map(lambda x: x, 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
observation = _model.Observation.from_dict(inputs)
start_time = time.monotonic()
outputs = {
"state": inputs["state"],
"actions": self._sample_actions(sample_rng_or_pytorch_device, observation, **sample_kwargs),
}
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)
outputs["policy_timing"] = {
"infer_ms": model_time * 1000,
}
return outputs
def infer_with_prefix_z(self, obs: dict, *, noise: np.ndarray | None = None) -> dict:
"""Run policy inference and return the COMET prefix latent used by A2C2.
This method preserves the normal output transforms for the action chunk,
and appends ``prefix_z`` after transforms so task-specific output
transforms cannot accidentally drop or reshape it.
"""
if self._is_pytorch_model or not self._supports_prefix_z:
raise NotImplementedError("infer_with_prefix_z is only supported for JAX policies with return_prefix_z.")
# Make a copy since transformations may modify the inputs in place.
inputs = jax.tree.map(lambda x: x, obs)
inputs = self._input_transform(inputs)
inputs = jax.tree.map(lambda x: jnp.asarray(x)[np.newaxis, ...], inputs)
self._rng, sample_rng = jax.random.split(self._rng)
sample_kwargs = dict(self._sample_kwargs)
if noise is not None:
noise = jnp.asarray(noise)
if noise.ndim == 2:
noise = noise[None, ...]
sample_kwargs["noise"] = noise
observation = _model.Observation.from_dict(inputs)
start_time = time.monotonic()
actions, prefix_z = self._sample_actions(
sample_rng,
observation,
**sample_kwargs,
return_prefix_z=True,
)
model_time = time.monotonic() - start_time
outputs = {
"state": inputs["state"],
"actions": actions,
}
outputs = jax.tree.map(lambda x: np.asarray(x[0, ...]), outputs)
prefix_z = np.asarray(prefix_z[0, ...])
outputs = self._output_transform(outputs)
outputs["prefix_z"] = prefix_z
outputs["policy_timing"] = {
"infer_ms": model_time * 1000,
}
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
|