Copybara import of the project:
-- 3a95b62f59e81bfef0f076afb173ecc14b27943d by Levi Burner <leviburner@gmail.com>: rollout prototype native threadpool for comparing to python threads -- efd8be1124ac839b902de45973a3ca8b9f2215e6 by Levi Burner <leviburner@gmail.com>: copy mjpcs threadpool into python bindings -- 75603eea3e8362e354a9675e8a6cd14e56ec3d28 by Levi Burner <leviburner@gmail.com>: rollout use threadpool as translation unit -- 06b90febd021663f6cc81fd7895e4d6e2008ed97 by Levi Burner <leviburner@gmail.com>: rollout add chunk_divisor parameter -- 298ab2f3c0d6e12530832c3cdbf784dd92d54806 by Levi Burner <leviburner@gmail.com>: rollout add native threading test -- 169cf9978e7abad6edd1392b8e6aab995e4f8f10 by Levi Burner <leviburner@gmail.com>: rollout exchange chunk_divisor arg for chunk_size -- 265af851d74432d261277d3dbda11cdef1841bc8 by Levi Burner <leviburner@gmail.com>: rollout fix cosmetics -- 1e8bffa88bf36190501b334bef31147e23db39f7 by Levi Burner <leviburner@gmail.com>: make native rollout a class instead of a function -- ba788214b047577f58c41ce0ab6c62c277cd8b0d by Levi Burner <leviburner@gmail.com>: rollout update docs and changelog -- e4cb7732319e04cba2ab2c2ad848c659f6309808 by Levi Burner <leviburner@gmail.com>: rollout don't register atexit handler for Rollout objects -- 5a08d2efdbbbb01d4b1231ff9a36a1dc44f4d9ee by Levi Burner <leviburner@gmail.com>: rollout nthread kwarg, rename shutdown_pool to close, fixups -- f622378543596a208339af0208fa3a70bf2a8007 by Levi Burner <leviburner@gmail.com>: rollout add missing .close() calls -- 50f3ebca43c53eac03f03943c34bb1e46967bd4f by Levi Burner <leviburner@gmail.com>: rollout return immediately COPYBARA_INTEGRATE_REVIEW=https://github.com/google-deepmind/mujoco/pull/2282 from aftersomemath:rollout-threaded 50f3ebca43c53eac03f03943c34bb1e46967bd4f PiperOrigin-RevId: 706744277 Change-Id: I1ab2263b7d6ce30cf1908aec8fd5f2eb976a19e6
This commit is contained in:
Levi Burner
committed by
Copybara-Service
Copybara-Service
parent
b26d6f0466
commit
a7eb6efd4e
+273
-128
@@ -14,6 +14,7 @@
|
||||
# ==============================================================================
|
||||
"""Roll out open-loop trajectories from initial states, get subsequent states and sensor values."""
|
||||
|
||||
import atexit
|
||||
from collections.abc import Sequence
|
||||
from typing import Optional, Union
|
||||
|
||||
@@ -23,9 +24,243 @@ import numpy as np
|
||||
from numpy import typing as npt
|
||||
|
||||
|
||||
class Rollout:
|
||||
"""Rollout object containing a thread pool for parallel rollouts."""
|
||||
|
||||
def __init__(self, *, nthread: Optional[int] = None):
|
||||
"""Construct a rollout object containing a thread pool for parallel rollouts.
|
||||
|
||||
Args:
|
||||
nthread: Number of threads in pool.
|
||||
If zero, this pool is not started and rollouts run on the calling thread.
|
||||
""" # fmt: skip
|
||||
self.nthread = 0 if nthread is None else nthread
|
||||
self.rollout_ = _rollout.Rollout(nthread=self.nthread)
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
self.close()
|
||||
|
||||
def close(self):
|
||||
del self.rollout_
|
||||
self.rollout_ = None
|
||||
|
||||
def rollout(
|
||||
self,
|
||||
model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
|
||||
data: Union[mujoco.MjData, Sequence[mujoco.MjData]],
|
||||
initial_state: npt.ArrayLike,
|
||||
control: Optional[npt.ArrayLike] = None,
|
||||
*, # require subsequent arguments to be named
|
||||
control_spec: int = mujoco.mjtState.mjSTATE_CTRL.value,
|
||||
skip_checks: bool = False,
|
||||
nstep: Optional[int] = None,
|
||||
initial_warmstart: Optional[npt.ArrayLike] = None,
|
||||
state: Optional[npt.ArrayLike] = None,
|
||||
sensordata: Optional[npt.ArrayLike] = None,
|
||||
chunk_size: Optional[int] = None,
|
||||
):
|
||||
"""Rolls out open-loop trajectories from initial states, get subsequent state and sensor values.
|
||||
|
||||
Python wrapper for rollout.cc, see documentation therein.
|
||||
Infers nroll and nstep.
|
||||
Tiles inputs with singleton dimensions.
|
||||
Allocates outputs if none are given.
|
||||
|
||||
Args:
|
||||
model: An instance or length nroll sequence of MjModel with the same size signature.
|
||||
data: Associated mjData instance or sequence of instances with length nthread.
|
||||
initial_state: Array of initial states from which to roll out trajectories.
|
||||
([nroll or 1] x nstate)
|
||||
control: Open-loop controls array to apply during the rollouts.
|
||||
([nroll or 1] x [nstep or 1] x ncontrol)
|
||||
control_spec: mjtState specification of control vectors.
|
||||
skip_checks: Whether to skip internal shape and type checks.
|
||||
nstep: Number of steps in rollouts (inferred if unspecified).
|
||||
initial_warmstart: Initial qfrc_warmstart array (optional).
|
||||
([nroll or 1] x nv)
|
||||
state: State output array (optional).
|
||||
(nroll x nstep x nstate)
|
||||
sensordata: Sensor data output array (optional).
|
||||
(nroll x nstep x nsensordata)
|
||||
chunk_size: Determines threadpool chunk size. If unspecified,
|
||||
chunk_size = max(1, nroll / (nthread * 10))
|
||||
|
||||
Returns:
|
||||
state:
|
||||
State output array, (nroll x nstep x nstate).
|
||||
sensordata:
|
||||
Sensor data output array, (nroll x nstep x nsensordata).
|
||||
|
||||
Raises:
|
||||
RuntimeError: rollout requested after thread pool shutdown.
|
||||
ValueError: bad shapes or sizes.
|
||||
""" # fmt: skip
|
||||
|
||||
if self.rollout_ is None:
|
||||
raise RuntimeError('rollout requested after thread pool shutdown')
|
||||
|
||||
# skip_checks shortcut:
|
||||
# don't infer nroll/nstep
|
||||
# don't support singleton expansion
|
||||
# don't allocate output arrays
|
||||
# just call rollout and return
|
||||
if skip_checks:
|
||||
self.rollout_.rollout(
|
||||
model,
|
||||
data,
|
||||
nstep,
|
||||
control_spec,
|
||||
initial_state,
|
||||
initial_warmstart,
|
||||
control,
|
||||
state,
|
||||
sensordata,
|
||||
chunk_size,
|
||||
)
|
||||
return state, sensordata
|
||||
|
||||
if not isinstance(model, mujoco.MjModel):
|
||||
model = list(model)
|
||||
|
||||
# check control_spec
|
||||
if control_spec & ~mujoco.mjtState.mjSTATE_USER.value:
|
||||
raise ValueError('control_spec can only contain bits in mjSTATE_USER')
|
||||
|
||||
# check types
|
||||
if nstep and not isinstance(nstep, int):
|
||||
raise ValueError('nstep must be an integer')
|
||||
if chunk_size and not isinstance(chunk_size, int):
|
||||
raise ValueError('chunk_size must be an integer')
|
||||
_check_must_be_numeric(
|
||||
initial_state=initial_state,
|
||||
initial_warmstart=initial_warmstart,
|
||||
control=control,
|
||||
state=state,
|
||||
sensordata=sensordata,
|
||||
)
|
||||
|
||||
# check number of dimensions
|
||||
_check_number_of_dimensions(
|
||||
2, initial_state=initial_state, initial_warmstart=initial_warmstart
|
||||
)
|
||||
_check_number_of_dimensions(
|
||||
3, control=control, state=state, sensordata=sensordata
|
||||
)
|
||||
|
||||
# ensure 2D, make contiguous, row-major (C ordering)
|
||||
initial_state = _ensure_2d(initial_state)
|
||||
initial_warmstart = _ensure_2d(initial_warmstart)
|
||||
|
||||
# ensure 3D, make contiguous, row-major (C ordering)
|
||||
control = _ensure_3d(control)
|
||||
state = _ensure_3d(state)
|
||||
sensordata = _ensure_3d(sensordata)
|
||||
|
||||
# infer nroll, check for incompatibilities
|
||||
nroll = _infer_dimension(
|
||||
0,
|
||||
1,
|
||||
initial_state=initial_state,
|
||||
initial_warmstart=initial_warmstart,
|
||||
control=control,
|
||||
state=state,
|
||||
sensordata=sensordata,
|
||||
)
|
||||
if isinstance(model, list) and nroll == 1:
|
||||
nroll = len(model)
|
||||
|
||||
if isinstance(model, list) and len(model) != nroll:
|
||||
raise ValueError(
|
||||
f'nroll inferred as {nroll} but model is length {len(model)}'
|
||||
)
|
||||
elif not isinstance(model, list):
|
||||
model = [model] # Use a length 1 list to simplify code below
|
||||
|
||||
if not isinstance(data, list):
|
||||
data = [data] # Use a length 1 list to simplify code below
|
||||
|
||||
# infer nstep, check for incompatibilities
|
||||
nstep = _infer_dimension(
|
||||
1, nstep or 1, control=control, state=state, sensordata=sensordata
|
||||
)
|
||||
|
||||
# get nstate/ncontrol/nv/nsensordata
|
||||
# check that they are equal across models
|
||||
nstate = mujoco.mj_stateSize(
|
||||
model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS.value
|
||||
)
|
||||
ncontrol = mujoco.mj_stateSize(model[0], control_spec)
|
||||
nv = model[0].nv
|
||||
nsensordata = model[0].nsensordata
|
||||
for m in model[1:]:
|
||||
if (
|
||||
nstate
|
||||
!= mujoco.mj_stateSize(m, mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
|
||||
or ncontrol != mujoco.mj_stateSize(m, control_spec)
|
||||
or nv != m.nv
|
||||
or nsensordata != m.nsensordata
|
||||
):
|
||||
raise ValueError('models are not compatible')
|
||||
|
||||
# check trailing dimensions
|
||||
_check_trailing_dimension(nstate, initial_state=initial_state, state=state)
|
||||
_check_trailing_dimension(ncontrol, control=control)
|
||||
_check_trailing_dimension(nv, initial_warmstart=initial_warmstart)
|
||||
_check_trailing_dimension(nsensordata, sensordata=sensordata)
|
||||
|
||||
# tile input arrays/lists if required (singleton expansion)
|
||||
model = model * nroll if len(model) == 1 else model
|
||||
initial_state = _tile_if_required(initial_state, nroll)
|
||||
initial_warmstart = _tile_if_required(initial_warmstart, nroll)
|
||||
control = _tile_if_required(control, nroll, nstep)
|
||||
|
||||
# allocate output if not provided
|
||||
if state is None:
|
||||
state = np.empty((nroll, nstep, nstate))
|
||||
if sensordata is None:
|
||||
sensordata = np.empty((nroll, nstep, nsensordata))
|
||||
|
||||
# call rollout
|
||||
self.rollout_.rollout(
|
||||
model,
|
||||
data,
|
||||
nstep,
|
||||
control_spec,
|
||||
initial_state,
|
||||
initial_warmstart,
|
||||
control,
|
||||
state,
|
||||
sensordata,
|
||||
chunk_size,
|
||||
)
|
||||
|
||||
# return outputs
|
||||
return state, sensordata
|
||||
|
||||
|
||||
persistent_rollout = None
|
||||
|
||||
|
||||
def shutdown_persistent_pool():
|
||||
"""Shutdown the persistent thread pool that is optionally created by rollout.
|
||||
|
||||
This is called automatically interpreter shutdown, but can also be called manually.
|
||||
""" # fmt: skip
|
||||
global persistent_rollout
|
||||
if persistent_rollout is not None:
|
||||
persistent_rollout.close()
|
||||
persistent_rollout = None
|
||||
|
||||
|
||||
atexit.register(shutdown_persistent_pool)
|
||||
|
||||
|
||||
def rollout(
|
||||
model: Union[mujoco.MjModel, Sequence[mujoco.MjModel]],
|
||||
data: mujoco.MjData,
|
||||
data: Union[mujoco.MjData, Sequence[mujoco.MjData]],
|
||||
initial_state: npt.ArrayLike,
|
||||
control: Optional[npt.ArrayLike] = None,
|
||||
*, # require subsequent arguments to be named
|
||||
@@ -35,6 +270,8 @@ def rollout(
|
||||
initial_warmstart: Optional[npt.ArrayLike] = None,
|
||||
state: Optional[npt.ArrayLike] = None,
|
||||
sensordata: Optional[npt.ArrayLike] = None,
|
||||
chunk_size: Optional[int] = None,
|
||||
persistent_pool: bool = False,
|
||||
):
|
||||
"""Rolls out open-loop trajectories from initial states, get subsequent states and sensor values.
|
||||
|
||||
@@ -44,8 +281,8 @@ def rollout(
|
||||
Allocates outputs if none are given.
|
||||
|
||||
Args:
|
||||
model: An mjModel or a sequence of MjModel with the same size signature.
|
||||
data: An associated mjData instance.
|
||||
model: An instance or length nroll sequence of MjModel with the same size signature.
|
||||
data: Associated mjData instance or sequence of instances with length nthread.
|
||||
initial_state: Array of initial states from which to roll out trajectories.
|
||||
([nroll or 1] x nstate)
|
||||
control: Open-loop controls array to apply during the rollouts.
|
||||
@@ -59,6 +296,9 @@ def rollout(
|
||||
(nroll x nstep x nstate)
|
||||
sensordata: Sensor data output array (optional).
|
||||
(nroll x nstep x nsensordata)
|
||||
chunk_size: Determines threadpool chunk size. If unspecified,
|
||||
chunk_size = max(1, nroll / (nthread * 10))
|
||||
persistent_pool: Determines if a persistent thread pool is created or reused.
|
||||
|
||||
Returns:
|
||||
state:
|
||||
@@ -69,136 +309,41 @@ def rollout(
|
||||
Raises:
|
||||
ValueError: bad shapes or sizes.
|
||||
""" # fmt: skip
|
||||
# skip_checks shortcut:
|
||||
# don't infer nroll/nstep
|
||||
# don't support singleton expansion
|
||||
# don't allocate output arrays
|
||||
# just call rollout and return
|
||||
if skip_checks:
|
||||
_rollout.rollout(
|
||||
if not isinstance(data, list):
|
||||
data = [data] # Use a length 1 list to simplify code below
|
||||
|
||||
nthread = len(data) if len(data) > 1 else 0
|
||||
|
||||
# Use a persistent thread pool if requested
|
||||
if persistent_pool:
|
||||
# Create or restart persistent threadpool
|
||||
global persistent_rollout
|
||||
if persistent_rollout is None:
|
||||
persistent_rollout = Rollout(nthread=nthread)
|
||||
if persistent_rollout.nthread != nthread:
|
||||
persistent_rollout.close()
|
||||
persistent_rollout = Rollout(nthread=nthread)
|
||||
rollout_ = persistent_rollout
|
||||
else:
|
||||
rollout_ = Rollout(nthread=nthread)
|
||||
|
||||
try:
|
||||
return rollout_.rollout(
|
||||
model,
|
||||
data,
|
||||
nstep,
|
||||
control_spec,
|
||||
initial_state,
|
||||
initial_warmstart,
|
||||
control,
|
||||
state,
|
||||
sensordata,
|
||||
control_spec=control_spec,
|
||||
skip_checks=skip_checks,
|
||||
nstep=nstep,
|
||||
initial_warmstart=initial_warmstart,
|
||||
state=state,
|
||||
sensordata=sensordata,
|
||||
chunk_size=chunk_size,
|
||||
)
|
||||
return state, sensordata
|
||||
|
||||
if not isinstance(model, mujoco.MjModel):
|
||||
model = list(model)
|
||||
|
||||
# check control_spec
|
||||
if control_spec & ~mujoco.mjtState.mjSTATE_USER.value:
|
||||
raise ValueError('control_spec can only contain bits in mjSTATE_USER')
|
||||
|
||||
# check types
|
||||
if nstep and not isinstance(nstep, int):
|
||||
raise ValueError('nstep must be an integer')
|
||||
_check_must_be_numeric(
|
||||
initial_state=initial_state,
|
||||
initial_warmstart=initial_warmstart,
|
||||
control=control,
|
||||
state=state,
|
||||
sensordata=sensordata,
|
||||
)
|
||||
|
||||
# check number of dimensions
|
||||
_check_number_of_dimensions(
|
||||
2, initial_state=initial_state, initial_warmstart=initial_warmstart
|
||||
)
|
||||
_check_number_of_dimensions(
|
||||
3, control=control, state=state, sensordata=sensordata
|
||||
)
|
||||
|
||||
# ensure 2D, make contiguous, row-major (C ordering)
|
||||
initial_state = _ensure_2d(initial_state)
|
||||
initial_warmstart = _ensure_2d(initial_warmstart)
|
||||
|
||||
# ensure 3D, make contiguous, row-major (C ordering)
|
||||
control = _ensure_3d(control)
|
||||
state = _ensure_3d(state)
|
||||
sensordata = _ensure_3d(sensordata)
|
||||
|
||||
# infer nroll, check for incompatibilities
|
||||
nroll = _infer_dimension(
|
||||
0,
|
||||
1,
|
||||
initial_state=initial_state,
|
||||
initial_warmstart=initial_warmstart,
|
||||
control=control,
|
||||
state=state,
|
||||
sensordata=sensordata,
|
||||
)
|
||||
if isinstance(model, list) and nroll == 1:
|
||||
nroll = len(model)
|
||||
|
||||
if isinstance(model, list) and len(model) != nroll:
|
||||
raise ValueError(
|
||||
f'nroll inferred as {nroll} but model is length {len(model)}'
|
||||
)
|
||||
elif not isinstance(model, list):
|
||||
model = [model] # Use a length 1 list to simplify code below
|
||||
|
||||
# infer nstep, check for incompatibilities
|
||||
nstep = _infer_dimension(
|
||||
1, nstep or 1, control=control, state=state, sensordata=sensordata
|
||||
)
|
||||
|
||||
# get nstate/ncontrol/nv/nsensordata
|
||||
# check that they are equal across models
|
||||
nstate = mujoco.mj_stateSize(
|
||||
model[0], mujoco.mjtState.mjSTATE_FULLPHYSICS.value
|
||||
)
|
||||
ncontrol = mujoco.mj_stateSize(model[0], control_spec)
|
||||
nv = model[0].nv
|
||||
nsensordata = model[0].nsensordata
|
||||
for m in model[1:]:
|
||||
if (
|
||||
nstate
|
||||
!= mujoco.mj_stateSize(m, mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
|
||||
or ncontrol != mujoco.mj_stateSize(m, control_spec)
|
||||
or nv != m.nv
|
||||
or nsensordata != m.nsensordata
|
||||
):
|
||||
raise ValueError('models are not compatible')
|
||||
|
||||
# check trailing dimensions
|
||||
_check_trailing_dimension(nstate, initial_state=initial_state, state=state)
|
||||
_check_trailing_dimension(ncontrol, control=control)
|
||||
_check_trailing_dimension(nv, initial_warmstart=initial_warmstart)
|
||||
_check_trailing_dimension(nsensordata, sensordata=sensordata)
|
||||
|
||||
# tile input arrays/lists if required (singleton expansion)
|
||||
model = model * nroll if len(model) == 1 else model
|
||||
initial_state = _tile_if_required(initial_state, nroll)
|
||||
initial_warmstart = _tile_if_required(initial_warmstart, nroll)
|
||||
control = _tile_if_required(control, nroll, nstep)
|
||||
|
||||
# allocate output if not provided
|
||||
if state is None:
|
||||
state = np.empty((nroll, nstep, nstate))
|
||||
if sensordata is None:
|
||||
sensordata = np.empty((nroll, nstep, nsensordata))
|
||||
|
||||
# call rollout
|
||||
_rollout.rollout(
|
||||
model,
|
||||
data,
|
||||
nstep,
|
||||
control_spec,
|
||||
initial_state,
|
||||
initial_warmstart,
|
||||
control,
|
||||
state,
|
||||
sensordata,
|
||||
)
|
||||
|
||||
# return outputs
|
||||
return state, sensordata
|
||||
finally:
|
||||
if not persistent_pool:
|
||||
rollout_.close()
|
||||
|
||||
|
||||
def _check_must_be_numeric(**kwargs):
|
||||
|
||||
Reference in New Issue
Block a user