| from typing import Dict |
|
|
| import numpy as np |
| import tree |
| from typing_extensions import override |
|
|
| from openpi_client import base_policy as _base_policy |
|
|
|
|
| class ActionChunkBroker(_base_policy.BasePolicy): |
| """Wraps a policy to return action chunks one-at-a-time. |
| |
| Assumes that the first dimension of all action fields is the chunk size. |
| |
| A new inference call to the inner policy is only made when the current |
| list of chunks is exhausted. |
| """ |
|
|
| def __init__(self, policy: _base_policy.BasePolicy, action_horizon: int): |
| self._policy = policy |
| self._action_horizon = action_horizon |
| self._cur_step: int = 0 |
|
|
| self._last_results: Dict[str, np.ndarray] | None = None |
|
|
| @override |
| def infer(self, obs: Dict) -> Dict: |
| if self._last_results is None: |
| self._last_results = self._policy.infer(obs) |
| self._cur_step = 0 |
|
|
| def slicer(x): |
| if isinstance(x, np.ndarray): |
| return x[self._cur_step, ...] |
| else: |
| return x |
|
|
| results = tree.map_structure(slicer, self._last_results) |
| self._cur_step += 1 |
|
|
| if self._cur_step >= self._action_horizon: |
| self._last_results = None |
|
|
| return results |
|
|
| @override |
| def reset(self) -> None: |
| self._policy.reset() |
| self._last_results = None |
| self._cur_step = 0 |
|
|
