File size: 4,456 Bytes
51c2c72 | 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 | 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
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,
):
self._sample_actions = nnx_utils.module_jit(model.sample_actions)
self._input_transform = _transforms.compose(transforms)
self._output_transform = _transforms.compose(output_transforms)
self._rng = rng or jax.random.key(0)
self._sample_kwargs = sample_kwargs or {}
self._metadata = dict(metadata or {})
# REPRO FIX: advertise the model's action shape so clients can size a per-query
# noise tensor without guessing (see infer()).
self._metadata.setdefault("action_horizon", int(getattr(model, "action_horizon", 0)))
self._metadata.setdefault("action_dim", int(getattr(model, "action_dim", 0)))
@override
def infer(self, obs: dict) -> dict: # type: ignore[misc]
# Make a copy since transformations may modify the inputs in place.
inputs = jax.tree.map(lambda x: x, obs)
# REPRO FIX: an optional client-supplied flow-matching noise tensor of shape
# [action_horizon, action_dim]. When present it fully determines the sample, so
# the result no longer depends on this server's cumulative RNG stream (which is
# advanced by every query from every client). Popped before the transforms.
client_noise = inputs.pop("noise", None) if isinstance(inputs, dict) else None
inputs = self._input_transform(inputs)
# Make a batch and convert to jax.Array.
inputs = jax.tree.map(lambda x: jnp.asarray(x)[np.newaxis, ...], inputs)
start_time = time.monotonic()
self._rng, sample_rng = jax.random.split(self._rng)
sample_kwargs = dict(self._sample_kwargs)
if client_noise is not None:
client_noise = np.asarray(client_noise, dtype=np.float32)
expected = (self._metadata["action_horizon"], self._metadata["action_dim"])
if tuple(client_noise.shape) != expected:
raise ValueError(f"client noise shape {client_noise.shape} != model (action_horizon, action_dim) {expected}")
sample_kwargs["noise"] = jnp.asarray(client_noise)[np.newaxis, ...]
outputs = {
"state": inputs["state"],
"actions": self._sample_actions(sample_rng, _model.Observation.from_dict(inputs), **sample_kwargs),
}
# Unbatch and convert to np.ndarray. # Unbatch and convert to np.ndarray.
outputs = jax.tree.map(lambda x: np.asarray(x[0, ...]), outputs)
model_time = time.monotonic() - start_time
outputs = self._output_transform(outputs)
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
|