mjrenderpool.py 文件源码

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