mujoco

# 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 typing import Optional

import mujoco
from mujoco import _rollout
import numpy as np
from numpy import typing as npt


def rollout(model: mujoco.MjModel,
            data: 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,
            nroll: Optional[int] = None,
            nstep: Optional[int] = None,
            initial_warmstart: Optional[npt.ArrayLike] = None,
            state: Optional[npt.ArrayLike] = None,
            sensordata: Optional[npt.ArrayLike] = None):
  """Rolls out open-loop trajectories from initial states, get subsequent states 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 mjModel instance.
    data: An associated mjData instance.
    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.
    nroll: Number of rollouts (inferred if unspecified).
    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)

  Returns:
    state:
      State output array, (nroll x nstep x nstate).
    sensordata:
      Sensor data output array, (nroll x nstep x nsensordata).

  Raises:
    ValueError: bad shapes or sizes.
  """
  # 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(model, data, nroll, nstep, control_spec, initial_state,
                     initial_warmstart, control, state, sensordata)
    return state, sensordata

  # 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 nroll and not isinstance(nroll, int):
    raise ValueError('nroll 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_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)

  # check trailing dimensions
  nstate = mujoco.mj_stateSize(model, mujoco.mjtState.mjSTATE_FULLPHYSICS.value)
  _check_trailing_dimension(nstate, initial_state=initial_state, state=state)
  ncontrol = mujoco.mj_stateSize(model, control_spec)
  _check_trailing_dimension(ncontrol, control=control)
  _check_trailing_dimension(model.nv, initial_warmstart=initial_warmstart)
  _check_trailing_dimension(model.nsensordata, sensordata=sensordata)

  # infer nroll, check for incompatibilities
  nroll = _infer_dimension(0, nroll or 1,
                           initial_state=initial_state,
                           initial_warmstart=initial_warmstart,
                           control=control,
                           state=state,
                           sensordata=sensordata)

  # infer nstep, check for incompatibilities
  nstep = _infer_dimension(1, nstep or 1,
                           control=control,
                           state=state,
                           sensordata=sensordata)

  # tile input arrays if required (singleton expansion)
  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, model.nsensordata))

  # call rollout
  _rollout.rollout(model, data, nroll, nstep, control_spec, initial_state,
                   initial_warmstart, control, state, sensordata)

  # return outputs
  return state, sensordata


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):
  """Infers dimension `dim` given guess `value` from set of arrays.

  Args:
    dim: Dimension to be inferred.
    value: Initial guess of inferred value (1: unknown).
    **kwargs: List of arrays which should all have the same size (or 1)
      along dimension dim.

  Returns:
    Inferred dimension.

  Raises:
    ValueError: If mismatch between array shapes or initial guess.
  """
  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} '
            f'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)