Version 2.1.2: Python bindings, OBJ assets support, bugfixes.
PiperOrigin-RevId: 434731612 Change-Id: I0cfda3e7a3d1c72036764986efc252ffa1b8c6b0
This commit is contained in:
Saran Tunyasuvunakool
@@ -0,0 +1,205 @@
|
||||
# Copyright 2022 DeepMind Technologies Limited
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
# ==============================================================================
|
||||
"""Roll out open-loop trajectories from initial states, get subsequent states and sensor values."""
|
||||
|
||||
from mujoco import _rollout
|
||||
import numpy as np
|
||||
|
||||
|
||||
def rollout(model, data, initial_state=None, ctrl=None,
|
||||
*, # require following arguments to be named
|
||||
skip_checks=False,
|
||||
nstate=None,
|
||||
nstep=None,
|
||||
initial_time=None,
|
||||
initial_warmstart=None,
|
||||
qfrc_applied=None,
|
||||
xfrc_applied=None,
|
||||
mocap=None,
|
||||
state=None,
|
||||
sensordata=None):
|
||||
"""Roll out open-loop trajectories from initial states, get subsequent states and sensor values.
|
||||
|
||||
This function serves as a Python wrapper for the C++ functionality in
|
||||
`rollout.cc`, please see documentation therein. This python funtion will
|
||||
infer `nstate` and `nstep`, tile input arguments with singleton dimensions,
|
||||
and allocate output arguments if none are given.
|
||||
"""
|
||||
# don't infer nstate/nstep, don't support singleton expansion, don't allocate
|
||||
# output arrays, just call rollout
|
||||
if skip_checks:
|
||||
_rollout.rollout(model, data, nstate, nstep, initial_state, initial_time,
|
||||
initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap,
|
||||
state, sensordata)
|
||||
return state, sensordata
|
||||
|
||||
# check types
|
||||
if nstate and not isinstance(nstate, int):
|
||||
raise ValueError('nstate must be an integer')
|
||||
if nstep and not isinstance(nstep, int):
|
||||
raise ValueError('nstep must be an integer')
|
||||
_check_must_be_numeric(
|
||||
initial_state=initial_state,
|
||||
initial_time=initial_time,
|
||||
initial_warmstart=initial_warmstart,
|
||||
ctrl=ctrl,
|
||||
qfrc_applied=qfrc_applied,
|
||||
xfrc_applied=xfrc_applied,
|
||||
mocap=mocap,
|
||||
state=state,
|
||||
sensordata=sensordata)
|
||||
|
||||
# check number of dimensions
|
||||
_check_number_of_dimensions(2,
|
||||
initial_state=initial_state,
|
||||
initial_time=initial_time,
|
||||
initial_warmstart=initial_warmstart)
|
||||
_check_number_of_dimensions(3,
|
||||
ctrl=ctrl,
|
||||
qfrc_applied=qfrc_applied,
|
||||
xfrc_applied=xfrc_applied,
|
||||
mocap=mocap,
|
||||
state=state,
|
||||
sensordata=sensordata)
|
||||
|
||||
# ensure 2D, make contiguous, row-major (C ordering)
|
||||
initial_state = _ensure_2d(initial_state)
|
||||
initial_time = _ensure_2d(initial_time)
|
||||
initial_warmstart = _ensure_2d(initial_warmstart)
|
||||
|
||||
# ensure 3D, make contiguous, row-major (C ordering)
|
||||
ctrl = _ensure_3d(ctrl)
|
||||
qfrc_applied = _ensure_3d(qfrc_applied)
|
||||
xfrc_applied = _ensure_3d(xfrc_applied)
|
||||
mocap = _ensure_3d(mocap)
|
||||
state = _ensure_3d(state)
|
||||
sensordata = _ensure_3d(sensordata)
|
||||
|
||||
# check trailing dimensions
|
||||
_check_trailing_dimension(model.nq + model.nv + model.na,
|
||||
initial_state=initial_state, state=state)
|
||||
_check_trailing_dimension(1, initial_time=initial_time)
|
||||
_check_trailing_dimension(model.nu, ctrl=ctrl)
|
||||
_check_trailing_dimension(model.nv, qfrc_applied=qfrc_applied)
|
||||
_check_trailing_dimension(model.nbody*6, xfrc_applied=xfrc_applied)
|
||||
_check_trailing_dimension(model.nmocap*7, mocap=mocap)
|
||||
_check_trailing_dimension(model.nsensordata, sensordata=sensordata)
|
||||
|
||||
# infer nstate, check for incompatibilities
|
||||
nstate = _infer_dimension(0, nstate or 1,
|
||||
initial_state=initial_state,
|
||||
initial_time=initial_time,
|
||||
initial_warmstart=initial_warmstart,
|
||||
ctrl=ctrl,
|
||||
qfrc_applied=qfrc_applied,
|
||||
xfrc_applied=xfrc_applied,
|
||||
mocap=mocap,
|
||||
state=state,
|
||||
sensordata=sensordata)
|
||||
|
||||
# infer nstep, check for incompatibilities
|
||||
nstep = _infer_dimension(1, nstep or 1,
|
||||
ctrl=ctrl,
|
||||
qfrc_applied=qfrc_applied,
|
||||
xfrc_applied=xfrc_applied,
|
||||
mocap=mocap,
|
||||
state=state,
|
||||
sensordata=sensordata)
|
||||
|
||||
# tile input arrays if required (singleton expansion)
|
||||
initial_state = _tile_if_required(initial_state, nstate)
|
||||
initial_time = _tile_if_required(initial_time, nstate)
|
||||
initial_warmstart = _tile_if_required(initial_warmstart, nstate)
|
||||
ctrl = _tile_if_required(ctrl, nstate, nstep)
|
||||
qfrc_applied = _tile_if_required(qfrc_applied, nstate, nstep)
|
||||
xfrc_applied = _tile_if_required(xfrc_applied, nstate, nstep)
|
||||
mocap = _tile_if_required(mocap, nstate, nstep)
|
||||
|
||||
# allocate output if not provided
|
||||
if state is None:
|
||||
state = np.empty((nstate, nstep, model.nq + model.nv + model.na))
|
||||
if sensordata is None:
|
||||
sensordata = np.empty((nstate, nstep, model.nsensordata))
|
||||
|
||||
# call rollout
|
||||
_rollout.rollout(model, data, nstate, nstep, initial_state, initial_time,
|
||||
initial_warmstart, ctrl, qfrc_applied, xfrc_applied, mocap,
|
||||
state, sensordata)
|
||||
|
||||
# return squeezed outputs
|
||||
return state.squeeze(), sensordata.squeeze()
|
||||
|
||||
def _check_must_be_numeric(**kwargs):
|
||||
for key, value in kwargs.items():
|
||||
if value is None:
|
||||
continue
|
||||
if not isinstance(value, np.ndarray) and not isinstance(value, float):
|
||||
raise ValueError(f'{key} must be a numpy array or float')
|
||||
|
||||
def _check_number_of_dimensions(ndim, **kwargs):
|
||||
for key, value in kwargs.items():
|
||||
if value is None:
|
||||
continue
|
||||
if value.ndim > ndim:
|
||||
raise ValueError(f'{key} can have at most {ndim} dimensions')
|
||||
|
||||
def _check_trailing_dimension(dim, **kwargs):
|
||||
for key, value in kwargs.items():
|
||||
if value is None:
|
||||
continue
|
||||
if value.shape[-1] != dim:
|
||||
raise ValueError(f'trailing dimension of {key} must be {dim}, got {value.shape[-1]}')
|
||||
|
||||
def _ensure_2d(arg):
|
||||
if arg is None:
|
||||
return None
|
||||
else:
|
||||
return np.ascontiguousarray(np.atleast_2d(arg), dtype=np.float64)
|
||||
|
||||
def _ensure_3d(arg):
|
||||
if arg is None:
|
||||
return None
|
||||
else:
|
||||
# np.atleast_3d adds both leading and trailing dims, we want only leading
|
||||
if arg.ndim == 0:
|
||||
arg = arg[np.newaxis, np.newaxis, np.newaxis, ...]
|
||||
elif arg.ndim == 1:
|
||||
arg = arg[np.newaxis, np.newaxis, ...]
|
||||
elif arg.ndim == 2:
|
||||
arg = arg[np.newaxis, ...]
|
||||
return np.ascontiguousarray(arg, dtype=np.float64)
|
||||
|
||||
def _infer_dimension(dim, value, **kwargs):
|
||||
for name, array in kwargs.items():
|
||||
if array is None:
|
||||
continue
|
||||
if array.shape[dim] != value:
|
||||
if value == 1:
|
||||
value = array.shape[dim]
|
||||
elif array.shape[dim] != 1:
|
||||
raise ValueError(
|
||||
f'dimension {dim} inferred as {value} but {name} has {array.shape[dim]}'
|
||||
)
|
||||
return value
|
||||
|
||||
def _tile_if_required(array, dim0, dim1=None):
|
||||
if array is None:
|
||||
return
|
||||
reps = np.ones(array.ndim, dtype=int)
|
||||
if array.shape[0] == 1:
|
||||
reps[0] = dim0
|
||||
if dim1 is not None and array.shape[1] == 1:
|
||||
reps[1] = dim1
|
||||
return np.tile(array, reps)
|
||||
Reference in New Issue
Block a user