mujoco_py/mjrenderpool.py

import ctypes
import inspect

from multiprocessing import Array, get_start_method, Pool, Value

import numpy as np


class RenderPoolStorage:
    """
    Helper object used for storing global data for worker processes.
    """

    __slots__ = ['shared_rgbs_array',
                 'shared_depths_array',
                 'device_id',
                 'sim',
                 'modder']


class MjRenderPool:
    """
    Utilizes a process pool to render a MuJoCo simulation across
    multiple GPU devices. This can scale the throughput linearly
    with the number of available GPUs. Throughput can also be
    slightly increased by using more than one worker per GPU.
    """

    DEFAULT_MAX_IMAGE_SIZE = 512 * 512  # in pixels

    def __init__(self, model, device_ids=1, n_workers=None,
                 max_batch_size=None, max_image_size=DEFAULT_MAX_IMAGE_SIZE,
                 modder=None):
        """
        Args:
        - model (PyMjModel): MuJoCo model to use for rendering
        - device_ids (int/list): list of device ids to use for rendering.
            One or more workers will be assigned to each device, depending
            on how many workers are requested.
        - n_workers (int): number of parallel processes in the pool. Defaults
            to the number of device ids.
        - max_batch_size (int): maximum number of states that can be rendered
            in batch using .render(). Defaults to the number of workers.
        - max_image_size (int): maximum number pixels in images requested
            by .render()
        - modder (Modder): modder to use for domain randomization.
        """
        self._closed, self.pool = False, None

        if not (modder is None or inspect.isclass(modder)):
            raise ValueError("modder must be a class")

        if isinstance(device_ids, int):
            device_ids = list(range(device_ids))
        else:
            assert isinstance(device_ids, list), (
                "device_ids must be list of integer")

        n_workers = n_workers or 1
        self._max_batch_size = max_batch_size or (len(device_ids) * n_workers)
        self._max_image_size = max_image_size

        array_size = self._max_image_size * self._max_batch_size

        self._shared_rgbs = Array(ctypes.c_uint8, array_size * 3)
        self._shared_depths = Array(ctypes.c_float, array_size)

        self._shared_rgbs_array = np.frombuffer(
            self._shared_rgbs.get_obj(), dtype=ctypes.c_uint8)
        assert self._shared_rgbs_array.size == (array_size * 3), (
            "Array size is %d, expected %d" % (
                self._shared_rgbs_array.size, array_size * 3))
        self._shared_depths_array = np.frombuffer(
            self._shared_depths.get_obj(), dtype=ctypes.c_float)
        assert self._shared_depths_array.size == array_size, (
            "Array size is %d, expected %d" % (
                self._shared_depths_array.size, array_size))

        worker_id = Value(ctypes.c_int)
        worker_id.value = 0

        if get_start_method() != "spawn":
            raise RuntimeError(
                "Start method must be set to 'spawn' for the "
                "render pool to work. That is, you must add the "
                "following to the _TOP_ of your main script, "
                "before any other imports (since they might be "
                "setting it otherwise):\n"
                "  import multiprocessing as mp\n"
                "  if __name__ == '__main__':\n"
                "    mp.set_start_method('spawn')\n")

        self.pool = Pool(
            processes=len(device_ids) * n_workers,
            initializer=MjRenderPool._worker_init,
            initargs=(
                model.get_mjb(),
                worker_id,
                device_ids,
                self._shared_rgbs,
                self._shared_depths,
                modder))

    @staticmethod
    def _worker_init(mjb_bytes, worker_id, device_ids,
                     shared_rgbs, shared_depths, modder):
        """
        Initializes the global state for the workers.
        """
        s = RenderPoolStorage()

        with worker_id.get_lock():
            proc_worker_id = worker_id.value
            worker_id.value += 1
        s.device_id = device_ids[proc_worker_id % len(device_ids)]

        s.shared_rgbs_array = np.frombuffer(
            shared_rgbs.get_obj(), dtype=ctypes.c_uint8)
        s.shared_depths_array = np.frombuffer(
            shared_depths.get_obj(), dtype=ctypes.c_float)

        # avoid a circular import
        from mujoco_py import load_model_from_mjb, MjRenderContext, MjSim
        s.sim = MjSim(load_model_from_mjb(mjb_bytes))
        # attach a render context to the sim (needs to happen before
        # modder is called, since it might need to upload textures
        # to the GPU).
        MjRenderContext(s.sim, device_id=s.device_id)

        if modder is not None:
            s.modder = modder(s.sim, random_state=proc_worker_id)
            s.modder.whiten_materials()
        else:
            s.modder = None

        global _render_pool_storage
        _render_pool_storage = s

    @staticmethod
    def _worker_render(worker_id, state, width, height,
                       camera_name, randomize):
        """
        Main target function for the workers.
        """
        s = _render_pool_storage

        forward = False
        if state is not None:
            s.sim.set_state(state)
            forward = True
        if randomize and s.modder is not None:
            s.modder.randomize()
            forward = True
        if forward:
            s.sim.forward()

        rgb_block = width * height * 3
        rgb_offset = rgb_block * worker_id
        rgb = s.shared_rgbs_array[rgb_offset:rgb_offset + rgb_block]
        rgb = rgb.reshape(height, width, 3)

        depth_block = width * height
        depth_offset = depth_block * worker_id
        depth = s.shared_depths_array[depth_offset:depth_offset + depth_block]
        depth = depth.reshape(height, width)

        rgb[:], depth[:] = s.sim.render(
            width, height, camera_name=camera_name, depth=True,
            device_id=s.device_id)

    def render(self, width, height, states=None, camera_name=None,
               depth=False, randomize=False, copy=True):
        """
        Renders the simulations in batch. If no states are provided,
        the max_batch_size will be used.

        Args:
        - width (int): width of image to render.
        - height (int): height of image to render.
        - states (list): list of MjSimStates; updates the states before
            rendering. Batch size will be number of states supplied.
        - camera_name (str): name of camera to render from.
        - depth (bool): if True, also return depth.
        - randomize (bool): calls modder.rand_all() before rendering.
        - copy (bool): return a copy rather than a reference

        Returns:
        - rgbs: NxHxWx3 numpy array of N images in batch of width W
            and height H.
        - depth: NxHxW numpy array of N images in batch of width W
            and height H. Only returned if depth=True.
        """
        if self._closed:
            raise RuntimeError("The pool has been closed.")

        if (width * height) > self._max_image_size:
            raise ValueError(
                "Requested image larger than maximum image size. Create "
                "a new RenderPool with a larger maximum image size.")
        if states is None:
            batch_size = self._max_batch_size
            states = [None] * batch_size
        else:
            batch_size = len(states)

        if batch_size > self._max_batch_size:
            raise ValueError(
                "Requested batch size larger than max batch size. Create "
                "a new RenderPool with a larger max batch size.")

        self.pool.starmap(
            MjRenderPool._worker_render,
            [(i, state, width, height, camera_name, randomize)
             for i, state in enumerate(states)])

        rgbs = self._shared_rgbs_array[:width * height * 3 * batch_size]
        rgbs = rgbs.reshape(batch_size, height, width, 3)
        if copy:
            rgbs = rgbs.copy()

        if depth:
            depths = self._shared_depths_array[:width * height * batch_size]
            depths = depths.reshape(batch_size, height, width).copy()
            if copy:
                depths = depths.copy()
            return rgbs, depths
        else:
            return rgbs

    def close(self):
        """
        Closes the pool and terminates child processes.
        """
        if not self._closed:
            if self.pool is not None:
                self.pool.close()
                self.pool.join()
            self._closed = True

    def __del__(self):
        self.close()