mujoco

// Copyright 2021 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.

#include <cstdio>
#include <cstdlib>
#include <cstring>

#include <mujoco/mujoco.h>

// select EGL, OSMESA or GLFW
#if defined(MJ_EGL)
  #include <EGL/egl.h>
#elif defined(MJ_OSMESA)
  #include <GL/osmesa.h>
OSMesaContext ctx;
unsigned char buffer[10000000];
#else
  #include <GLFW/glfw3.h>
#endif

#include "array_safety.h"
namespace mju = ::mujoco::sample_util;

//-------------------------------- global data -----------------------------------------------------

// MuJoCo model and data
mjModel* m = 0;
mjData*  d = 0;

// MuJoCo visualization
mjvScene   scn;
mjvCamera  cam;
mjvOption  opt;
mjrContext con;


//-------------------------------- utility functions -----------------------------------------------

// load model, init simulation and rendering
void initMuJoCo(const char* filename) {
  // load and compile
  char error[1000] = "Could not load binary model";
  if (std::strlen(filename) > 4 && !std::strcmp(filename + std::strlen(filename) - 4, ".mjb")) {
    m = mj_loadModel(filename, 0);
  } else {
    m = mj_loadXML(filename, 0, error, 1000);
  }
  if (!m) { mju_error("Load model error: %s", error); }

  // make data, run one computation to initialize all fields
  d = mj_makeData(m);
  mj_forward(m, d);

  // initialize visualization data structures
  mjv_defaultCamera(&cam);
  mjv_defaultOption(&opt);
  mjv_defaultScene(&scn);
  mjr_defaultContext(&con);

  // create scene and context
  mjv_makeScene(m, &scn, 2000);
  mjr_makeContext(m, &con, 200);

  // default free camera
  mjv_defaultFreeCamera(m, &cam);
}


// deallocate everything
void closeMuJoCo(void) {
  mj_deleteData(d);
  mj_deleteModel(m);
  mjr_freeContext(&con);
  mjv_freeScene(&scn);
}


// create OpenGL context/window
void initOpenGL(void) {
  //------------------------ EGL
#if defined(MJ_EGL)
  // desired config
  // clang-format off
  const EGLint configAttribs[] = {
    EGL_RED_SIZE,           8,
    EGL_GREEN_SIZE,         8,
    EGL_BLUE_SIZE,          8,
    EGL_ALPHA_SIZE,         8,
    EGL_DEPTH_SIZE,         24,
    EGL_STENCIL_SIZE,       8,
    EGL_COLOR_BUFFER_TYPE,  EGL_RGB_BUFFER,
    EGL_SURFACE_TYPE,       EGL_PBUFFER_BIT,
    EGL_RENDERABLE_TYPE,    EGL_OPENGL_BIT,
    EGL_NONE
  };
  // clang-format on

  // get default display
  EGLDisplay eglDpy = eglGetDisplay(EGL_DEFAULT_DISPLAY);
  if (eglDpy == EGL_NO_DISPLAY) {
    mju_error("Could not get EGL display, error 0x%x\n", eglGetError());
  }

  // initialize
  EGLint major, minor;
  if (eglInitialize(eglDpy, &major, &minor) != EGL_TRUE) {
    mju_error("Could not initialize EGL, error 0x%x\n", eglGetError());
  }

  // choose config
  EGLint    numConfigs;
  EGLConfig eglCfg;
  if (eglChooseConfig(eglDpy, configAttribs, &eglCfg, 1, &numConfigs) != EGL_TRUE) {
    mju_error("Could not choose EGL config, error 0x%x\n", eglGetError());
  }

  // bind OpenGL API
  if (eglBindAPI(EGL_OPENGL_API) != EGL_TRUE) {
    mju_error("Could not bind EGL OpenGL API, error 0x%x\n", eglGetError());
  }

  // create context
  EGLContext eglCtx = eglCreateContext(eglDpy, eglCfg, EGL_NO_CONTEXT, NULL);
  if (eglCtx == EGL_NO_CONTEXT) {
    mju_error("Could not create EGL context, error 0x%x\n", eglGetError());
  }

  // make context current, no surface (let OpenGL handle FBO)
  if (eglMakeCurrent(eglDpy, EGL_NO_SURFACE, EGL_NO_SURFACE, eglCtx) != EGL_TRUE) {
    mju_error("Could not make EGL context current, error 0x%x\n", eglGetError());
  }

  //------------------------ OSMESA
#elif defined(MJ_OSMESA)
  // create context
  ctx = OSMesaCreateContextExt(GL_RGBA, 24, 8, 8, 0);
  if (!ctx) { mju_error("OSMesa context creation failed"); }

  // make current
  if (!OSMesaMakeCurrent(ctx, buffer, GL_UNSIGNED_BYTE, 800, 800)) {
    mju_error("OSMesa make current failed");
  }

  //------------------------ GLFW
#else
  // init GLFW
  if (!glfwInit()) { mju_error("Could not initialize GLFW"); }

  // create invisible window, single-buffered
  glfwWindowHint(GLFW_VISIBLE, 0);
  glfwWindowHint(GLFW_DOUBLEBUFFER, GLFW_FALSE);
  GLFWwindow* window = glfwCreateWindow(800, 800, "Invisible window", NULL, NULL);
  if (!window) { mju_error("Could not create GLFW window"); }

  // make context current
  glfwMakeContextCurrent(window);
#endif
}


// close OpenGL context/window
void closeOpenGL(void) {
  //------------------------ EGL
#if defined(MJ_EGL)
  // get current display
  EGLDisplay eglDpy = eglGetCurrentDisplay();
  if (eglDpy == EGL_NO_DISPLAY) { return; }

  // get current context
  EGLContext eglCtx = eglGetCurrentContext();

  // release context
  eglMakeCurrent(eglDpy, EGL_NO_SURFACE, EGL_NO_SURFACE, EGL_NO_CONTEXT);

  // destroy context if valid
  if (eglCtx != EGL_NO_CONTEXT) { eglDestroyContext(eglDpy, eglCtx); }

  // terminate display
  eglTerminate(eglDpy);

  //------------------------ OSMESA
#elif defined(MJ_OSMESA)
  OSMesaDestroyContext(ctx);

  //------------------------ GLFW
#else
  // terminate GLFW (crashes with Linux NVidia drivers)
  #if defined(__APPLE__) || defined(_WIN32)
  glfwTerminate();
  #endif
#endif
}


//-------------------------------- main function ---------------------------------------------------

int main(int argc, const char** argv) {
  // check command-line arguments
  if (argc < 5 || argc > 6) {
    std::printf(" USAGE:  record modelfile duration fps rgbfile [adddepth]\n");
    return EXIT_FAILURE;
  }

  // parse numeric arguments
  double duration = 10, fps = 30;
  std::sscanf(argv[2], "%lf", &duration);
  std::sscanf(argv[3], "%lf", &fps);

  // initialize OpenGL and MuJoCo
  initOpenGL();
  initMuJoCo(argv[1]);

  // set rendering to offscreen buffer
  mjr_setBuffer(mjFB_OFFSCREEN, &con);
  if (con.currentBuffer != mjFB_OFFSCREEN) {
    std::printf("Warning: offscreen rendering not supported, using default/window framebuffer\n");
  }

  // get size of active renderbuffer
  mjrRect viewport = mjr_maxViewport(&con);
  int     W        = viewport.width;
  int     H        = viewport.height;

  // allocate rgb and depth buffers
  unsigned char* rgb   = (unsigned char*)std::malloc(3 * W * H);
  float*         depth = (float*)std::malloc(sizeof(float) * W * H);
  if (!rgb || !depth) { mju_error("Could not allocate buffers"); }

  // create output rgb file
  std::FILE* fp = std::fopen(argv[4], "wb");
  if (!fp) { mju_error("Could not open rgbfile for writing"); }

  int adddepth = 1;
  if (argc > 5 && std::sscanf(argv[5], "%d", &adddepth) != 1) {
    mju_error("Invalid adddepth argument");
  }

  // main loop
  double frametime  = 0;
  int    framecount = 0;
  while (d->time < duration) {
    // render new frame if it is time (or first frame)
    if ((d->time - frametime) > 1 / fps || frametime == 0) {
      // update abstract scene
      mjv_updateScene(m, d, &opt, NULL, &cam, mjCAT_ALL, &scn);

      // render scene in offscreen buffer
      mjr_render(viewport, &scn, &con);

      // add time stamp in upper-left corner
      char stamp[50];
      mju::sprintf_arr(stamp, "Time = %.3f", d->time);
      mjr_overlay(mjFONT_NORMAL, mjGRID_TOPLEFT, viewport, stamp, NULL, &con);

      // read rgb and depth buffers
      mjr_readPixels(rgb, depth, viewport, &con);

      // insert subsampled depth image in lower-left corner of rgb image
      if (adddepth) {
        const int NS = 3;  // depth image subsampling
        for (int r = 0; r < H; r += NS) {
          for (int c = 0; c < W; c += NS) {
            int adr      = (r / NS) * W + c / NS;
            rgb[3 * adr] = rgb[3 * adr + 1] = rgb[3 * adr + 2] =
                (unsigned char)((1.0f - depth[r * W + c]) * 255.0f);
          }
        }
      }

      // write rgb image to file
      std::fwrite(rgb, 3, W * H, fp);

      // print every 10 frames: '.' if ok, 'x' if OpenGL error
      if (((framecount++) % 10) == 0) {
        if (mjr_getError()) {
          std::printf("x");
        } else {
          std::printf(".");
        }
      }

      // save simulation time
      frametime = d->time;
    }

    // advance simulation
    mj_step(m, d);
  }
  std::printf("\n");

  // close file, free buffers
  std::fclose(fp);
  std::free(rgb);
  std::free(depth);

  // close MuJoCo and OpenGL
  closeMuJoCo();
  closeOpenGL();

  return EXIT_SUCCESS;
}