From 695892ba1f7f062208d15ca95f8e7740acccef69 Mon Sep 17 00:00:00 2001 From: Jonathan Raiman Date: Sat, 18 Jun 2022 14:07:59 -0700 Subject: [PATCH 1/8] partial support for newer mujoco --- robogym/envs/dactyl/tests/test_locked.py | 7 +- robogym/envs/rearrange/common/utils.py | 4 +- .../holdouts/tests/test_stability.py | 2 +- robogym/envs/rearrange/simulation/base.py | 2 +- robogym/mujoco/callbacks.cpp | 312 ++++++++++++ robogym/mujoco/constants.py | 461 ++++++++++++++++++ robogym/mujoco/helpers.py | 11 +- robogym/mujoco/mujoco_xml.py | 36 +- robogym/mujoco/simulation_interface.py | 46 +- robogym/mujoco/test/test_mujoco_utils.py | 32 +- robogym/mujoco/warning_buffer.py | 15 +- robogym/randomization/env.py | 7 +- robogym/randomization/sim.py | 2 +- .../shadow_hand/mujoco/mujoco_shadow_hand.py | 2 +- robogym/viewer/env_viewer.py | 9 +- robogym/viewer/holdout_creation_viewer.py | 2 +- robogym/viewer/robot_control_viewer.py | 2 +- setup.py | 43 +- 18 files changed, 904 insertions(+), 91 deletions(-) create mode 100644 robogym/mujoco/callbacks.cpp diff --git a/robogym/envs/dactyl/tests/test_locked.py b/robogym/envs/dactyl/tests/test_locked.py index 92193c6..ec2437b 100644 --- a/robogym/envs/dactyl/tests/test_locked.py +++ b/robogym/envs/dactyl/tests/test_locked.py @@ -1,5 +1,4 @@ import numpy as np -from mujoco_py import ignore_mujoco_warnings from numpy.testing import assert_allclose from robogym.envs.dactyl.common.cube_utils import on_palm @@ -50,9 +49,9 @@ def test_locked_cube(): ) assert env.unwrapped.sim.model.joint_names == expected_joints - with ignore_mujoco_warnings(): - for _ in range(20): - obs, _, _, _ = env.step(env.action_space.nvec // 2) + # with ignore_mujoco_warnings(): + for _ in range(20): + obs, _, _, _ = env.step(env.action_space.nvec // 2) is_on_palm.append(on_palm(env.unwrapped.sim)) diff --git a/robogym/envs/rearrange/common/utils.py b/robogym/envs/rearrange/common/utils.py index 502b7d5..17ec721 100644 --- a/robogym/envs/rearrange/common/utils.py +++ b/robogym/envs/rearrange/common/utils.py @@ -12,14 +12,14 @@ import numpy as np import trimesh from collision import Poly, Vector, collide -from mujoco_py import MjSim, const from numpy.random import RandomState -from robogym.mujoco.mujoco_xml import ASSETS_DIR, MujocoXML +from robogym.mujoco.mujoco_xml import MjSim, ASSETS_DIR, MujocoXML from robogym.utils.env_utils import InvalidSimulationError from robogym.utils.mesh import get_vertices_bounding_box, subdivide_mesh from robogym.utils.misc import robogym_path from robogym.utils.rotation import mat2quat, quat2mat, quat_conjugate, uniform_quat +from robogym.mujoco import constants as const MATERIAL_DIR = robogym_path("envs", "rearrange", "materials") diff --git a/robogym/envs/rearrange/holdouts/tests/test_stability.py b/robogym/envs/rearrange/holdouts/tests/test_stability.py index c037621..3b93e5d 100644 --- a/robogym/envs/rearrange/holdouts/tests/test_stability.py +++ b/robogym/envs/rearrange/holdouts/tests/test_stability.py @@ -2,7 +2,7 @@ from typing import List import numpy as np -from mujoco_py import const +from robogym.mujoco import constants as const from robogym.robot_env import RobotEnv from robogym.utils.env_utils import load_env diff --git a/robogym/envs/rearrange/simulation/base.py b/robogym/envs/rearrange/simulation/base.py index 5cf4a4f..e0f1dcf 100644 --- a/robogym/envs/rearrange/simulation/base.py +++ b/robogym/envs/rearrange/simulation/base.py @@ -5,7 +5,7 @@ import attr import numpy as np from gym.envs.robotics import rotations -from mujoco_py import const +from robogym.mujoco import constants as const from robogym.envs.rearrange.common.utils import ( PlacementArea, diff --git a/robogym/mujoco/callbacks.cpp b/robogym/mujoco/callbacks.cpp new file mode 100644 index 0000000..a034243 --- /dev/null +++ b/robogym/mujoco/callbacks.cpp @@ -0,0 +1,312 @@ +#include +#include +#include +#include "Python.h" + +namespace { + enum USER_DEFINED_ACTUATOR_DATA { + IDX_CONTROLLER_TYPE = 0, + NUM_ACTUATOR_DATA = 1 + }; + enum USER_DEFINED_ACTUATOR_PARAMS { + IDX_PROPORTIONAL_GAIN = 0, + IDX_INTEGRAL_TIME_CONSTANT = 1, + IDX_INTEGRAL_MAX_CLAMP = 2, + IDX_DERIVATIVE_TIME_CONSTANT = 3, + IDX_DERIVATIVE_GAIN_SMOOTHING = 4, + IDX_ERROR_DEADBAND = 5 + }; + enum USER_DEFINED_CONTROLLER_DATA_PID { + IDX_INTEGRAL_ERROR = 0, + IDX_ERROR = 1, + IDX_DERIVATIVE_ERROR = 2, + NUM_USER_DATA_PER_ACT_PID = 3 + }; + + + enum USER_DEFINED_ACTUATOR_PARAMS_CASCADE { + IDX_CAS_PROPORTIONAL_GAIN = 0, + IDX_CAS_INTEGRAL_TIME_CONSTANT = 1, + IDX_CAS_INTEGRAL_MAX_CLAMP = 2, + IDX_CAS_DERIVATIVE_TIME_CONSTANT = 3, + IDX_CAS_DERIVATIVE_GAIN_SMOOTHING = 4, + IDX_CAS_PROPORTIONAL_GAIN_V = 5, + IDX_CAS_INTEGRAL_TIME_CONSTANT_V = 6, + IDX_CAS_INTEGRAL_MAX_CLAMP_V = 7, + IDX_CAS_EMA_SMOOTH = 8, + IDX_CAS_MAX_VEL = 9 + }; + + enum USER_DEFINED_CONTROLLER_DATA_CASCADE { + IDX_CAS_ERROR = 0, + IDX_CAS_INTEGRAL_ERROR = 1, + IDX_CAS_DERIVATIVE_ERROR = 2, + IDX_CAS_INTEGRAL_ERROR_V = 3, + IDX_CAS_STORED_EMA_SMOOTH = 4, + NUM_USER_DATA_PER_ACT_CAS = 5 + }; + const int CONTROLLER_TYPE_PI_CASCADE = 1; + const int NUM_USER_DATA_PER_ACT = std::max(int(NUM_USER_DATA_PER_ACT_PID), int(NUM_USER_DATA_PER_ACT_CAS)); + struct PIDErrors { + double error; + double integral_error; + double derivative_error; + PIDErrors(double error_, double integral_error_, double derivative_error_) : error(error_), integral_error(integral_error_), derivative_error(derivative_error_) {} + }; + + struct PIDOutput { + PIDErrors errors; + double output; + PIDOutput(PIDErrors errors_, double output_) : errors(errors_), output(output_) {} + }; + + struct PIDParameters { + double dt_seconds; // PID sampling time. + double setpoint; + double feedback; + double Kp; + double error_deadband; + double integral_max_clamp; + double integral_time_const; + double derivative_gain_smoothing; + double derivative_time_const; + PIDErrors previous_errors; + PIDParameters( + double dt_seconds_, + double setpoint_, + double feedback_, + double Kp_, + double error_deadband_, + double integral_max_clamp_, + double integral_time_const_, + double derivative_gain_smoothing_, + double derivative_time_const_, + PIDErrors previous_errors_) : dt_seconds(dt_seconds_), + setpoint(setpoint_), + feedback(feedback_), + Kp(Kp_), + error_deadband(error_deadband_), + integral_max_clamp(integral_max_clamp_), + integral_time_const(integral_time_const_), + derivative_gain_smoothing(derivative_gain_smoothing_), + derivative_time_const(derivative_time_const_), + previous_errors(previous_errors_) {} + }; +} + + +PIDOutput _pid(PIDParameters parameters) { + // A general purpose PID controller implemented in the standard form-> + // Kp == Kp + // Ki == Kp/Ti + // Kd == Kp*Td + + // In this situation, Kp is a knob to tune the agressiveness, wheras Ti and Td will + // change the response time of the system in a predictable way. Lower Ti or Td means + // that the system will respond to error more quickly/agressively. + + // error deadband: if set will shrink error within to 0.0 + // clamp on integral term: helps on saturation problem in I. + // derivative smoothing term: reduces high frequency noise in D-> + + // :param parameters: PID parameters + // :return: A PID output struct containing the control output and the error state + + double error = parameters.setpoint - parameters.feedback; + // Clamp error that's within the error deadband + if (fabs(error) < parameters.error_deadband) { + error = 0.0; + } + + // Compute derivative error + double derivative_error = (error - parameters.previous_errors.error) / parameters.dt_seconds; + + derivative_error = (1.0 - parameters.derivative_gain_smoothing) * parameters.previous_errors.derivative_error + \ + parameters.derivative_gain_smoothing * derivative_error; + + double derivative_error_term = derivative_error * parameters.derivative_time_const; + + // Update and clamp integral error + double integral_error = parameters.previous_errors.integral_error; + integral_error += error * parameters.dt_seconds; + integral_error = fmax(-parameters.integral_max_clamp, fmin(parameters.integral_max_clamp, integral_error)); + + double integral_error_term = 0.0; + if (parameters.integral_time_const != 0) { + integral_error_term = integral_error / parameters.integral_time_const; + } + double f = parameters.Kp * (error + integral_error_term + derivative_error_term); + return PIDOutput(PIDErrors(error, integral_error, derivative_error), f); +} + + +extern "C" int get_NUM_USER_DATA_PER_ACT() {return NUM_USER_DATA_PER_ACT;} +extern "C" mjtNum c_zero_gains(const mjModel* m, const mjData* d, int id) {return 0.0;} + +extern "C" mjtNum c_pi_cascade_bias(const mjModel* m, const mjData* d, int id) { + // A cascaded PID-PI controller implementation that can control position + // and velocity setpoints for a given actuator. + + // The cascaded controller is implemented as a nested position-velocity controller. A PID loop is + // wrapped around desired position and a PI loop is wrapped around desired velocity. An exponential + // moving average filter is applied to the commanded position input (d->ctrl). Additionally the controller + // is gravity compensated via the `qfrc_bias` term-> + // The PID-PI gains are set as part of gainprm in the following order: + // `gainprm=" Kp_pos -> Position loop proportional gain + // Ti_pos -> Position loop integral time constant + // Ti_max_clamp_pos -> Position loop integral error clamp + // Td_pos -> Position loop derivative time constant + // Td_smooth_pos -> Position loop derivative smoothing (EMA) + // Kp_vel -> Velocity loop proportional gain + // Ti_vel -> Velocity loop integral time constant + // max_clamp_vel -> Velocity loop integral error clamp + // ema_smooth_factor -> Exponential moving average (EMA) on desired position + // max_vel"` -> Clamped velocity limit (applied in positive and negative direction) + double dt_in_sec = m->opt.timestep; + int NGAIN = 10.0; + double smooth_pos_setpoint;; + + double Kp_cas = m->actuator_gainprm[id * NGAIN + IDX_CAS_PROPORTIONAL_GAIN]; + + // Apply Exponential Moving Average smoothing to the position setpoint, applying warmstart + // on the first iteration. + if (d->time > 0.0) { + auto ctrl_ema = d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_CAS_STORED_EMA_SMOOTH]; + auto smoothing_factor = m->actuator_gainprm[id * NGAIN + IDX_CAS_EMA_SMOOTH]; + smooth_pos_setpoint = (smoothing_factor * ctrl_ema) + (1 - smoothing_factor) * d->ctrl[id]; + } else { + smooth_pos_setpoint = d->ctrl[id]; + } + d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_CAS_STORED_EMA_SMOOTH] = smooth_pos_setpoint; + double des_vel; + + if (Kp_cas != 0) { + // Run a position PID loop and use the result to set a desired velocity signal + auto pos_output = _pid(PIDParameters( + /*dt_seconds=*/dt_in_sec, + /*setpoint=*/smooth_pos_setpoint, + /*feedback=*/d->actuator_length[id], + /*Kp=*/Kp_cas, + /*error_deadband=*/0.0, + /*integral_max_clamp=*/m->actuator_gainprm[id * NGAIN + IDX_CAS_INTEGRAL_MAX_CLAMP], + /*integral_time_const=*/m->actuator_gainprm[id * NGAIN + IDX_CAS_INTEGRAL_TIME_CONSTANT], + /*derivative_gain_smoothing=*/m->actuator_gainprm[id * NGAIN + IDX_CAS_DERIVATIVE_GAIN_SMOOTHING], + /*derivative_time_const=*/m->actuator_gainprm[id * NGAIN + IDX_CAS_DERIVATIVE_TIME_CONSTANT], + /*previous_errors=*/PIDErrors( + /*error=*/d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_CAS_ERROR], + /*integral_error=*/d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_CAS_DERIVATIVE_ERROR], + /*derivative_error=*/d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_CAS_INTEGRAL_ERROR] + ))); + + // Save errors + d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_CAS_ERROR] = pos_output.errors.error; + d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_CAS_DERIVATIVE_ERROR] = pos_output.errors.derivative_error; + d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_CAS_INTEGRAL_ERROR] = pos_output.errors.integral_error; + des_vel = pos_output.output; + } else { + // If P gain on position loop is zero, only use the velocity controller + des_vel = d->ctrl[id]; + } + + // Clamp max angular velocity + auto max_qvel = m->actuator_gainprm[id * NGAIN + IDX_CAS_MAX_VEL]; + des_vel = fmax(-max_qvel, fmin(max_qvel, des_vel)); + + auto vel_output = _pid(PIDParameters( + /*dt_seconds=*/dt_in_sec, + /*setpoint=*/des_vel, + /*feedback=*/d->actuator_velocity[id], + /*Kp=*/m->actuator_gainprm[id * NGAIN + IDX_CAS_PROPORTIONAL_GAIN_V], + /*error_deadband=*/0.0, + /*integral_max_clamp=*/m->actuator_gainprm[id * NGAIN + IDX_CAS_INTEGRAL_MAX_CLAMP_V], + /*integral_time_const=*/m->actuator_gainprm[id * NGAIN + IDX_CAS_INTEGRAL_TIME_CONSTANT_V], + /*derivative_gain_smoothing=*/0.0, + /*derivative_time_const=*/0.0, + /*previous_errors=*/PIDErrors( + /*error=*/0.0, + /*integral_error=*/0.0, + /*derivative_error=*/d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_CAS_INTEGRAL_ERROR_V] + ))); + + // Save errors + d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_CAS_INTEGRAL_ERROR_V] = vel_output.errors.integral_error; + + // Limit max torque at the output + double effort_limit_low = m->actuator_forcerange[id * 2]; + double effort_limit_high = m->actuator_forcerange[id * 2 + 1]; + + auto f = vel_output.output; + + // Gravity compensation + f += d->qfrc_bias[id]; + + if (effort_limit_low != 0.0 or effort_limit_high != 0.0) { + f = fmax(effort_limit_low, fmin(effort_limit_high, f)); + } + return f; +} + +extern "C" mjtNum c_pid_bias(const mjModel* m, const mjData* d, int id) { + // To activate PID, set gainprm="Kp Ti Td iClamp errBand iSmooth" in a general type actuator in mujoco xml + double dt_in_sec = m->opt.timestep; + int NGAIN = 10.0; + auto result = _pid(PIDParameters( + /*dt_seconds=*/dt_in_sec, + /*setpoint=*/d->ctrl[id], + /*feedback=*/d->actuator_length[id], + /*Kp=*/m->actuator_gainprm[id * NGAIN + IDX_PROPORTIONAL_GAIN], + /*error_deadband=*/m->actuator_gainprm[id * NGAIN + IDX_ERROR_DEADBAND], + /*integral_max_clamp=*/m->actuator_gainprm[id * NGAIN + IDX_INTEGRAL_MAX_CLAMP], + /*integral_time_const=*/m->actuator_gainprm[id * NGAIN + IDX_INTEGRAL_TIME_CONSTANT], + /*derivative_gain_smoothing=*/m->actuator_gainprm[id * NGAIN + IDX_DERIVATIVE_GAIN_SMOOTHING], + /*derivative_time_const=*/m->actuator_gainprm[id * NGAIN + IDX_DERIVATIVE_TIME_CONSTANT], + /*previous_errors=*/PIDErrors( + /*error=*/d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_ERROR], + /*integral_error=*/d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_DERIVATIVE_ERROR], + /*derivative_error=*/d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_INTEGRAL_ERROR] + ))); + + d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_ERROR] = result.errors.error; + d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_DERIVATIVE_ERROR] = result.errors.derivative_error; + d->userdata[id * NUM_USER_DATA_PER_ACT + IDX_INTEGRAL_ERROR] = result.errors.integral_error; + mjtNum f = result.output; + + double effort_limit_low = m->actuator_forcerange[id * 2]; + double effort_limit_high = m->actuator_forcerange[id * 2 + 1]; + + if ((effort_limit_low != 0.0) || (effort_limit_high != 0.0)) { + f = fmax(effort_limit_low, fmin(effort_limit_high, f)); + } + return f; +} + +extern "C" mjtNum c_custom_bias(const mjModel*m, const mjData*d, int id) { + // Switches between PID and Cascaded PID-PI type custom bias computation based on the + // defined actuator's actuator_user field-> + // user="1": Cascade PID-PI + // default: PID + // :param m: mjModel + // :param d: mjData + // :param id: actuator ID + // :return: Custom actuator force + int controller_type = int(m->actuator_user[id * m->nuser_actuator + IDX_CONTROLLER_TYPE]); + + if (controller_type == CONTROLLER_TYPE_PI_CASCADE) { + return c_pi_cascade_bias(m, d, id); + } + return c_pid_bias(m, d, id); +} + + +static struct PyModuleDef callbacks_definition = { + PyModuleDef_HEAD_INIT, + "callbacks", + "extension code for supporting mujoco callbacks.", + -1, + NULL +}; + +PyMODINIT_FUNC PyInit_callbacks(void) { + Py_Initialize(); + return PyModule_Create(&callbacks_definition); +} \ No newline at end of file diff --git a/robogym/mujoco/constants.py b/robogym/mujoco/constants.py index e3408c3..c6ac2d4 100644 --- a/robogym/mujoco/constants.py +++ b/robogym/mujoco/constants.py @@ -59,3 +59,464 @@ class MujocoEquality(Enum): mjEQ_JOINT = 2 # couple the values of two scalar joints with cubic mjEQ_TENDON = 3 # couple the lengths of two tendons with cubic mjEQ_DISTANCE = 4 # fix the contact distance betweent two geoms + +# Automatically generated. Do not modify! + +###### const from defines ###### +MINVAL = 1e-15 # minimum value in any denominator +PI = 3.141592653589793 # +MAXVAL = 10000000000.0 # maximum value in qpos, qvel, qacc +MINMU = 1e-05 # minimum friction coefficient +MINIMP = 0.0001 # minimum constraint impedance +MAXIMP = 0.9999 # maximum constraint impedance +MAXCONPAIR = 50.0 # maximum number of contacts per geom pair +MAXVFS = 2000.0 # maximum number of files in virtual file system +MAXVFSNAME = 1000.0 # maximum filename size in virtual file system +NEQDATA = 7.0 # number of eq_data fields +NDYN = 10.0 # number of actuator dynamics parameters +NGAIN = 10.0 # number of actuator gain parameters +NBIAS = 10.0 # number of actuator bias parameters +NREF = 2.0 # number of solver reference parameters +NIMP = 5.0 # number of solver impedance parameters +NSOLVER = 1000.0 # size of mjData.solver_XXX arrays +NGROUP = 6.0 # number of geom, site, joint groups with visflags +MAXOVERLAY = 500.0 # maximum number of characters in overlay text +MAXLINE = 100.0 # maximum number of lines per plot +MAXLINEPNT = 1000.0 # maximum number points per line +MAXPLANEGRID = 200.0 # maximum number of grid divisions for plane +NAUX = 10.0 # number of auxiliary buffers +MAXTEXTURE = 1000.0 # maximum number of textures +MAXUISECT = 10.0 # maximum number of sections +MAXUIITEM = 80.0 # maximum number of items per section +MAXUITEXT = 300.0 # maximum number of chars in edittext and other +MAXUINAME = 40.0 # maximum number of chars in name +MAXUIMULTI = 35.0 # maximum number of radio/select items in group +MAXUIEDIT = 7.0 # maximum number of elements in edit list +MAXUIRECT = 25.0 # maximum number of rectangles +SEPCLOSED = 1000.0 # closed state of adjustable separator +KEY_ESCAPE = 256.0 # +KEY_ENTER = 257.0 # +KEY_TAB = 258.0 # +KEY_BACKSPACE = 259.0 # +KEY_INSERT = 260.0 # +KEY_DELETE = 261.0 # +KEY_RIGHT = 262.0 # +KEY_LEFT = 263.0 # +KEY_DOWN = 264.0 # +KEY_UP = 265.0 # +KEY_PAGE_UP = 266.0 # +KEY_PAGE_DOWN = 267.0 # +KEY_HOME = 268.0 # +KEY_END = 269.0 # +KEY_F1 = 290.0 # +KEY_F2 = 291.0 # +KEY_F3 = 292.0 # +KEY_F4 = 293.0 # +KEY_F5 = 294.0 # +KEY_F6 = 295.0 # +KEY_F7 = 296.0 # +KEY_F8 = 297.0 # +KEY_F9 = 298.0 # +KEY_F10 = 299.0 # +KEY_F11 = 300.0 # +KEY_F12 = 301.0 # + +###### const from enums ###### + + # _mjtDisableBit +DSBL_CONSTRAINT = 1 +DSBL_EQUALITY = 2 +DSBL_FRICTIONLOSS = 4 +DSBL_LIMIT = 8 +DSBL_CONTACT = 16 +DSBL_PASSIVE = 32 +DSBL_GRAVITY = 64 +DSBL_CLAMPCTRL = 128 +DSBL_WARMSTART = 256 +DSBL_FILTERPARENT = 512 +DSBL_ACTUATION = 1024 +DSBL_REFSAFE = 2048 +NDISABLE = 12 + + # _mjtEnableBit +ENBL_OVERRIDE = 1 +ENBL_ENERGY = 2 +ENBL_FWDINV = 4 +ENBL_SENSORNOISE = 8 +NENABLE = 4 + + # _mjtJoint +JNT_FREE = 0 +JNT_BALL = 1 +JNT_SLIDE = 2 +JNT_HINGE = 3 + + # _mjtGeom +GEOM_PLANE = 0 +GEOM_HFIELD = 1 +GEOM_SPHERE = 2 +GEOM_CAPSULE = 3 +GEOM_ELLIPSOID = 4 +GEOM_CYLINDER = 5 +GEOM_BOX = 6 +GEOM_MESH = 7 +NGEOMTYPES = 8 +GEOM_ARROW = 100 +GEOM_ARROW1 = 101 +GEOM_ARROW2 = 102 +GEOM_LINE = 103 +GEOM_SKIN = 104 +GEOM_LABEL = 105 +GEOM_NONE = 1001 + + # _mjtCamLight +CAMLIGHT_FIXED = 0 +CAMLIGHT_TRACK = 1 +CAMLIGHT_TRACKCOM = 2 +CAMLIGHT_TARGETBODY = 3 +CAMLIGHT_TARGETBODYCOM = 4 + + # _mjtTexture +TEXTURE_2D = 0 +TEXTURE_CUBE = 1 +TEXTURE_SKYBOX = 2 + + # _mjtIntegrator +INT_EULER = 0 +INT_RK4 = 1 + + # _mjtCollision +COL_ALL = 0 +COL_PAIR = 1 +COL_DYNAMIC = 2 + + # _mjtCone +CONE_PYRAMIDAL = 0 +CONE_ELLIPTIC = 1 + + # _mjtJacobian +JAC_DENSE = 0 +JAC_SPARSE = 1 +JAC_AUTO = 2 + + # _mjtSolver +SOL_PGS = 0 +SOL_CG = 1 +SOL_NEWTON = 2 + + # _mjtEq +EQ_CONNECT = 0 +EQ_WELD = 1 +EQ_JOINT = 2 +EQ_TENDON = 3 +EQ_DISTANCE = 4 + + # _mjtWrap +WRAP_NONE = 0 +WRAP_JOINT = 1 +WRAP_PULLEY = 2 +WRAP_SITE = 3 +WRAP_SPHERE = 4 +WRAP_CYLINDER = 5 + + # _mjtTrn +TRN_JOINT = 0 +TRN_JOINTINPARENT = 1 +TRN_SLIDERCRANK = 2 +TRN_TENDON = 3 +TRN_SITE = 4 +TRN_UNDEFINED = 1000 + + # _mjtDyn +DYN_NONE = 0 +DYN_INTEGRATOR = 1 +DYN_FILTER = 2 +DYN_MUSCLE = 3 +DYN_USER = 4 + + # _mjtGain +GAIN_FIXED = 0 +GAIN_MUSCLE = 1 +GAIN_USER = 2 + + # _mjtBias +BIAS_NONE = 0 +BIAS_AFFINE = 1 +BIAS_MUSCLE = 2 +BIAS_USER = 3 + + # _mjtObj +OBJ_UNKNOWN = 0 +OBJ_BODY = 1 +OBJ_XBODY = 2 +OBJ_JOINT = 3 +OBJ_DOF = 4 +OBJ_GEOM = 5 +OBJ_SITE = 6 +OBJ_CAMERA = 7 +OBJ_LIGHT = 8 +OBJ_MESH = 9 +OBJ_SKIN = 10 +OBJ_HFIELD = 11 +OBJ_TEXTURE = 12 +OBJ_MATERIAL = 13 +OBJ_PAIR = 14 +OBJ_EXCLUDE = 15 +OBJ_EQUALITY = 16 +OBJ_TENDON = 17 +OBJ_ACTUATOR = 18 +OBJ_SENSOR = 19 +OBJ_NUMERIC = 20 +OBJ_TEXT = 21 +OBJ_TUPLE = 22 +OBJ_KEY = 23 + + # _mjtConstraint +CNSTR_EQUALITY = 0 +CNSTR_FRICTION_DOF = 1 +CNSTR_FRICTION_TENDON = 2 +CNSTR_LIMIT_JOINT = 3 +CNSTR_LIMIT_TENDON = 4 +CNSTR_CONTACT_FRICTIONLESS = 5 +CNSTR_CONTACT_PYRAMIDAL = 6 +CNSTR_CONTACT_ELLIPTIC = 7 + + # _mjtConstraintState +CNSTRSTATE_SATISFIED = 0 +CNSTRSTATE_QUADRATIC = 1 +CNSTRSTATE_LINEARNEG = 2 +CNSTRSTATE_LINEARPOS = 3 +CNSTRSTATE_CONE = 4 + + # _mjtSensor +SENS_TOUCH = 0 +SENS_ACCELEROMETER = 1 +SENS_VELOCIMETER = 2 +SENS_GYRO = 3 +SENS_FORCE = 4 +SENS_TORQUE = 5 +SENS_MAGNETOMETER = 6 +SENS_RANGEFINDER = 7 +SENS_JOINTPOS = 8 +SENS_JOINTVEL = 9 +SENS_TENDONPOS = 10 +SENS_TENDONVEL = 11 +SENS_ACTUATORPOS = 12 +SENS_ACTUATORVEL = 13 +SENS_ACTUATORFRC = 14 +SENS_BALLQUAT = 15 +SENS_BALLANGVEL = 16 +SENS_JOINTLIMITPOS = 17 +SENS_JOINTLIMITVEL = 18 +SENS_JOINTLIMITFRC = 19 +SENS_TENDONLIMITPOS = 20 +SENS_TENDONLIMITVEL = 21 +SENS_TENDONLIMITFRC = 22 +SENS_FRAMEPOS = 23 +SENS_FRAMEQUAT = 24 +SENS_FRAMEXAXIS = 25 +SENS_FRAMEYAXIS = 26 +SENS_FRAMEZAXIS = 27 +SENS_FRAMELINVEL = 28 +SENS_FRAMEANGVEL = 29 +SENS_FRAMELINACC = 30 +SENS_FRAMEANGACC = 31 +SENS_SUBTREECOM = 32 +SENS_SUBTREELINVEL = 33 +SENS_SUBTREEANGMOM = 34 +SENS_USER = 35 + + # _mjtStage +STAGE_NONE = 0 +STAGE_POS = 1 +STAGE_VEL = 2 +STAGE_ACC = 3 + + # _mjtDataType +DATATYPE_REAL = 0 +DATATYPE_POSITIVE = 1 +DATATYPE_AXIS = 2 +DATATYPE_QUATERNION = 3 + + # _mjtLRMode +LRMODE_NONE = 0 +LRMODE_MUSCLE = 1 +LRMODE_MUSCLEUSER = 2 +LRMODE_ALL = 3 + + # _mjtWarning +WARN_INERTIA = 0 +WARN_CONTACTFULL = 1 +WARN_CNSTRFULL = 2 +WARN_VGEOMFULL = 3 +WARN_BADQPOS = 4 +WARN_BADQVEL = 5 +WARN_BADQACC = 6 +WARN_BADCTRL = 7 +NWARNING = 8 + + # _mjtTimer +TIMER_STEP = 0 +TIMER_FORWARD = 1 +TIMER_INVERSE = 2 +TIMER_POSITION = 3 +TIMER_VELOCITY = 4 +TIMER_ACTUATION = 5 +TIMER_ACCELERATION = 6 +TIMER_CONSTRAINT = 7 +TIMER_POS_KINEMATICS = 8 +TIMER_POS_INERTIA = 9 +TIMER_POS_COLLISION = 10 +TIMER_POS_MAKE = 11 +TIMER_POS_PROJECT = 12 +NTIMER = 13 + + # _mjtCatBit +CAT_STATIC = 1 +CAT_DYNAMIC = 2 +CAT_DECOR = 4 +CAT_ALL = 7 + + # _mjtMouse +MOUSE_NONE = 0 +MOUSE_ROTATE_V = 1 +MOUSE_ROTATE_H = 2 +MOUSE_MOVE_V = 3 +MOUSE_MOVE_H = 4 +MOUSE_ZOOM = 5 +MOUSE_SELECT = 6 + + # _mjtPertBit +PERT_TRANSLATE = 1 +PERT_ROTATE = 2 + + # _mjtCamera +CAMERA_FREE = 0 +CAMERA_TRACKING = 1 +CAMERA_FIXED = 2 +CAMERA_USER = 3 + + # _mjtLabel +LABEL_NONE = 0 +LABEL_BODY = 1 +LABEL_JOINT = 2 +LABEL_GEOM = 3 +LABEL_SITE = 4 +LABEL_CAMERA = 5 +LABEL_LIGHT = 6 +LABEL_TENDON = 7 +LABEL_ACTUATOR = 8 +LABEL_CONSTRAINT = 9 +LABEL_SKIN = 10 +LABEL_SELECTION = 11 +LABEL_SELPNT = 12 +LABEL_CONTACTFORCE = 13 +NLABEL = 14 + + # _mjtFrame +FRAME_NONE = 0 +FRAME_BODY = 1 +FRAME_GEOM = 2 +FRAME_SITE = 3 +FRAME_CAMERA = 4 +FRAME_LIGHT = 5 +FRAME_WORLD = 6 +NFRAME = 7 + + # _mjtVisFlag +VIS_CONVEXHULL = 0 +VIS_TEXTURE = 1 +VIS_JOINT = 2 +VIS_ACTUATOR = 3 +VIS_CAMERA = 4 +VIS_LIGHT = 5 +VIS_TENDON = 6 +VIS_RANGEFINDER = 7 +VIS_CONSTRAINT = 8 +VIS_INERTIA = 9 +VIS_SCLINERTIA = 10 +VIS_PERTFORCE = 11 +VIS_PERTOBJ = 12 +VIS_CONTACTPOINT = 13 +VIS_CONTACTFORCE = 14 +VIS_CONTACTSPLIT = 15 +VIS_TRANSPARENT = 16 +VIS_AUTOCONNECT = 17 +VIS_COM = 18 +VIS_SELECT = 19 +VIS_STATIC = 20 +VIS_SKIN = 21 +NVISFLAG = 22 + + # _mjtRndFlag +RND_SHADOW = 0 +RND_WIREFRAME = 1 +RND_REFLECTION = 2 +RND_ADDITIVE = 3 +RND_SKYBOX = 4 +RND_FOG = 5 +RND_HAZE = 6 +RND_SEGMENT = 7 +RND_IDCOLOR = 8 +NRNDFLAG = 9 + + # _mjtStereo +STEREO_NONE = 0 +STEREO_QUADBUFFERED = 1 +STEREO_SIDEBYSIDE = 2 + + # _mjtGridPos +GRID_TOPLEFT = 0 +GRID_TOPRIGHT = 1 +GRID_BOTTOMLEFT = 2 +GRID_BOTTOMRIGHT = 3 + + # _mjtFramebuffer +FB_WINDOW = 0 +FB_OFFSCREEN = 1 + + # _mjtFontScale +FONTSCALE_50 = 50 +FONTSCALE_100 = 100 +FONTSCALE_150 = 150 +FONTSCALE_200 = 200 +FONTSCALE_250 = 250 +FONTSCALE_300 = 300 + + # _mjtFont +FONT_NORMAL = 0 +FONT_SHADOW = 1 +FONT_BIG = 2 + + # _mjtButton +BUTTON_NONE = 0 +BUTTON_LEFT = 1 +BUTTON_RIGHT = 2 +BUTTON_MIDDLE = 3 + + # _mjtEvent +EVENT_NONE = 0 +EVENT_MOVE = 1 +EVENT_PRESS = 2 +EVENT_RELEASE = 3 +EVENT_SCROLL = 4 +EVENT_KEY = 5 +EVENT_RESIZE = 6 + + # _mjtItem +ITEM_END = -2 +ITEM_SECTION = -1 +ITEM_SEPARATOR = 0 +ITEM_STATIC = 1 +ITEM_BUTTON = 2 +ITEM_CHECKINT = 3 +ITEM_CHECKBYTE = 4 +ITEM_RADIO = 5 +ITEM_RADIOLINE = 6 +ITEM_SELECT = 7 +ITEM_SLIDERINT = 8 +ITEM_SLIDERNUM = 9 +ITEM_EDITINT = 10 +ITEM_EDITNUM = 11 +ITEM_EDITTXT = 12 +NITEM = 13 diff --git a/robogym/mujoco/helpers.py b/robogym/mujoco/helpers.py index 7bcc5a7..69b13de 100644 --- a/robogym/mujoco/helpers.py +++ b/robogym/mujoco/helpers.py @@ -1,8 +1,7 @@ import itertools import typing -import mujoco_py -import mujoco_py.generated.const +from robogym.mujoco import constants as const def joint_qpos_ids(model, joint_name: str) -> typing.List[int]: @@ -40,13 +39,13 @@ def joint_qvel_ids_from_prefix(model, joint_prefix): def joint_type_name(joint_type: int) -> str: - if joint_type == mujoco_py.generated.const.JNT_FREE: + if joint_type == const.JNT_FREE: return "free" - if joint_type == mujoco_py.generated.const.JNT_BALL: + if joint_type == const.JNT_BALL: return "ball" - if joint_type == mujoco_py.generated.const.JNT_SLIDE: + if joint_type == const.JNT_SLIDE: return "slide" - if joint_type == mujoco_py.generated.const.JNT_HINGE: + if joint_type == const.JNT_HINGE: return "hinge" raise AssertionError(f"unsupported joint type: {joint_type}") diff --git a/robogym/mujoco/mujoco_xml.py b/robogym/mujoco/mujoco_xml.py index de34ad8..481d657 100644 --- a/robogym/mujoco/mujoco_xml.py +++ b/robogym/mujoco/mujoco_xml.py @@ -2,7 +2,7 @@ import typing import xml.etree.ElementTree as et -import mujoco_py +import mujoco import numpy as np ASSETS_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "../assets")) @@ -31,28 +31,17 @@ class StaleMjSimError(Exception): """ Exception indicating the MjSim instance is stale and should no longer be used. """ - pass -class MjSim(mujoco_py.MjSim): - """ - There are environments e.g. rearrange environment which recreates - sim after reach env reset. This can cause potential bugs caused by - other components still caching instance of old sim. These bugs are usually - quite tricky to find. This class makes it easier to find these bugs by allowing - invalidating the sim instance so any access to properties of stale sim instance - will cause error. - """ - - __slots__ = ("_stale", "_xml") - - def __init__(self, model, **kwargs): - # Note: we don't need to call super.__init__ because MjSim use __cinit__ - # for initialization which happens automatically before subclass __init__ - # is called. +class MjSim(object): + def __init__(self, model, data, nsubsteps=1): self._stale: bool = False - self._xml = model.get_xml() + self._model = model + self._data = data + self._xml = None# model.get_xml() + self._nsubsteps = nsubsteps + print(f"nsubsteps = {nsubsteps}") def get_xml(self): """ @@ -76,12 +65,12 @@ def is_stale(self): @property def data(self): self._ensure_not_stale() - return super().data + return self._data @property def model(self): self._ensure_not_stale() - return super().model + return self._model def _ensure_not_stale(self): if self._stale: @@ -256,8 +245,9 @@ def build(self, output_filename=None, meshdir=None, texturedir=None, **kwargs): with open(output_filename, "wt") as f: f.write(xml_string) - mj_model = mujoco_py.load_model_from_xml(xml_string) - return MjSim(mj_model, **kwargs) + mj_model = mujoco.MjModel.from_xml_string(xml_string) + mj_data = mujoco.MjData(mj_model) + return MjSim(mj_model, mj_data, **kwargs) ############################################################################################### # MODIFICATIONS diff --git a/robogym/mujoco/simulation_interface.py b/robogym/mujoco/simulation_interface.py index 2e095d3..c2f8278 100644 --- a/robogym/mujoco/simulation_interface.py +++ b/robogym/mujoco/simulation_interface.py @@ -2,8 +2,8 @@ from typing import Dict, List import attr -from mujoco_py import MjSimState, cymj - +import ctypes +import mujoco from robogym.mujoco.helpers import ( joint_qpos_ids, joint_qpos_ids_from_prefix, @@ -11,6 +11,19 @@ joint_qvel_ids_from_prefix, ) from robogym.mujoco.mujoco_xml import MjSim +from . import callbacks + +_callbacks = ctypes.PyDLL(callbacks.__file__) + +c_zero_gains = _callbacks.c_zero_gains +c_custom_bias = _callbacks.c_custom_bias +_callbacks.get_NUM_USER_DATA_PER_ACT.restype = ctypes.c_int + +NUM_USER_DATA_PER_ACT = _callbacks.get_NUM_USER_DATA_PER_ACT() +IDX_CONTROLLER_TYPE = 0 +NUM_ACTUATOR_DATA = 1 + + @attr.s(auto_attribs=True) @@ -74,18 +87,31 @@ def mujoco_viewer(self): """ Get a nicely-interactive version of the mujoco viewer """ - if self._mujoco_viewer is None: - # Inline import since this is only relevant on platforms - # which have GLFW support. - from mujoco_py.mjviewer import MjViewer # noqa + raise NotImplementedError() + # if self._mujoco_viewer is None: + # # Inline import since this is only relevant on platforms + # # which have GLFW support. + # from mujoco_py.mjviewer import MjViewer # noqa - self._mujoco_viewer = MjViewer(self.sim) + # self._mujoco_viewer = MjViewer(self.sim) return self._mujoco_viewer def enable_pid(self): """ Enable our custom PID controller code for the actuators with 'user' type """ - cymj.set_pid_control(self.sim.model, self.sim.data) + m = self.sim.model + d = self.sim.data + + if m.nuserdata < m.nu * NUM_USER_DATA_PER_ACT: + raise Exception('nuserdata is not set large enough to store PID internal states.') + + if m.nuser_actuator < m.nu * NUM_ACTUATOR_DATA: + raise Exception('nuser_actuator is not set large enough to store controller types') + + for i in range(m.nuserdata): + d.userdata[i] = 0.0 + mujoco.set_mjcb_act_gain(c_zero_gains) + mujoco.set_mjcb_act_bias(c_custom_bias) ############################################################################################### # SUBCLASS REGISTRATION @@ -151,7 +177,7 @@ def qvel(self): """ Returns copy of full sim qvel. """ return self.sim.data.qvel.copy() - def get_state(self) -> MjSimState: + def get_state(self): return self.sim.get_state() ############################################################################################### @@ -168,7 +194,7 @@ def add_qpos(self, group_name, value): """ Sets qpos for a given group. """ self.sim.data.qpos[self.qpos_idxs[group_name]] += value - def set_state(self, state: MjSimState): + def set_state(self, state): self.sim.set_state(state) ############################################################################################### diff --git a/robogym/mujoco/test/test_mujoco_utils.py b/robogym/mujoco/test/test_mujoco_utils.py index 34419c8..0994165 100644 --- a/robogym/mujoco/test/test_mujoco_utils.py +++ b/robogym/mujoco/test/test_mujoco_utils.py @@ -1,7 +1,6 @@ import random import numpy as np -from mujoco_py import cymj, functions from numpy.random.mtrand import _rand as global_randstate from robogym.mujoco.forward_kinematics import ForwardKinematics @@ -196,23 +195,22 @@ def test_remove_elem(): assert ref_xml.strip() == ball_without_joint.xml_string().strip() -def test_mj_error_callback(): - message = None - called = False +# def test_mj_error_callback(): +# message = None +# called = False - def callback(msg): - nonlocal message - message = msg.decode() - raise RuntimeError(message) +# def callback(msg): +# nonlocal message +# message = msg.decode() +# raise RuntimeError(message) - cymj.set_error_callback(callback) +# cymj.set_error_callback(callback) - try: - with cymj.wrap_mujoco_warning(): - functions.mju_error("error") - except RuntimeError as e: - assert e.args[0] == "error" - assert message == "error" - called = True +# try: +# with cymj.wrap_mujoco_warning(): +# except RuntimeError as e: +# assert e.args[0] == "error" +# assert message == "error" +# called = True - assert called +# assert called diff --git a/robogym/mujoco/warning_buffer.py b/robogym/mujoco/warning_buffer.py index baf91f3..aee0bea 100644 --- a/robogym/mujoco/warning_buffer.py +++ b/robogym/mujoco/warning_buffer.py @@ -1,8 +1,5 @@ import collections import logging - -import mujoco_py.cymj as cymj - logger = logging.getLogger(__name__) @@ -21,7 +18,8 @@ def error_callback(message): # Set it once for all the processes -cymj.set_error_callback(error_callback) +# TODO(jonathan)?? +# cymj.set_error_callback(error_callback) class MjWarningBuffer: @@ -55,10 +53,11 @@ def warnings(self): def enter(self): """ Enable collecting warnings """ - if self._prev_user_callback is None: - self._prev_user_callback = cymj.get_warning_callback() + # if self._prev_user_callback is None: + # self._prev_user_callback = cymj.get_warning_callback() - cymj.set_warning_callback(self._intercept_warning) + # cymj.set_warning_callback(self._intercept_warning) + pass def clear(self): """ Reset warning buffer """ @@ -67,7 +66,7 @@ def clear(self): def exit(self): """ Stop collecting warnings """ if self._prev_user_callback is not None: - cymj.set_warning_callback(self._prev_user_callback) + # cymj.set_warning_callback(self._prev_user_callback) self._prev_user_callback = None def __enter__(self): diff --git a/robogym/randomization/env.py b/robogym/randomization/env.py index ff44be6..dae9551 100644 --- a/robogym/randomization/env.py +++ b/robogym/randomization/env.py @@ -78,16 +78,11 @@ def value_in_range(_, attribute, value): ) -class _RandomizableParam(NamedTuple, Generic[VType]): - +class _RandomizableParam(NamedTuple): name: str - value_type: Type[VType] - default: VType - value_range: Tuple[VType, VType] - parent_instance: Any diff --git a/robogym/randomization/sim.py b/robogym/randomization/sim.py index 4a785ac..867e3c7 100644 --- a/robogym/randomization/sim.py +++ b/robogym/randomization/sim.py @@ -3,7 +3,6 @@ from typing import List, Union import numpy as np -from mujoco_py import MjSim from numpy.random import RandomState from robogym.mujoco.constants import OPT_FIELDS, PID_GAIN_PARAMS @@ -12,6 +11,7 @@ FloatRandomizerParameter, RandomizerParameter, ) +from robogym.mujoco.mujoco_xml import MjSim from robogym.utils.rotation import random_unity2 diff --git a/robogym/robot/shadow_hand/mujoco/mujoco_shadow_hand.py b/robogym/robot/shadow_hand/mujoco/mujoco_shadow_hand.py index 9952f46..9503a74 100644 --- a/robogym/robot/shadow_hand/mujoco/mujoco_shadow_hand.py +++ b/robogym/robot/shadow_hand/mujoco/mujoco_shadow_hand.py @@ -1,5 +1,5 @@ import numpy as np -from mujoco_py.generated import const +from robogym.mujoco import constants as const from robogym.mujoco.simulation_interface import SimulationInterface from robogym.robot.shadow_hand.hand_forward_kinematics import ( diff --git a/robogym/viewer/env_viewer.py b/robogym/viewer/env_viewer.py index 395fbd2..7980a23 100644 --- a/robogym/viewer/env_viewer.py +++ b/robogym/viewer/env_viewer.py @@ -5,7 +5,8 @@ import glfw import numpy as np from gym.spaces import Box, MultiDiscrete -from mujoco_py import MjViewer, const, ignore_mujoco_warnings +import robogym.mujoco.constants as const +from mujoco_py import MjViewer from robogym.utils.misc import pretty @@ -124,10 +125,8 @@ def run(self, once=False): self.process_events() self.update_sim(self.env.unwrapped.sim) self.add_extra_menu() - - with ignore_mujoco_warnings(): - obs, reward, done, info = self._run_step(self._get_action()) - self.last_step_result = deepcopy((obs, reward, done, info)) + obs, reward, done, info = self._run_step(self._get_action()) + self.last_step_result = deepcopy((obs, reward, done, info)) self.add_overlay(const.GRID_BOTTOMRIGHT, "done", str(done)) self.render() diff --git a/robogym/viewer/holdout_creation_viewer.py b/robogym/viewer/holdout_creation_viewer.py index df9f015..dcdfacc 100644 --- a/robogym/viewer/holdout_creation_viewer.py +++ b/robogym/viewer/holdout_creation_viewer.py @@ -4,7 +4,7 @@ import glfw import numpy as np -from mujoco_py import const +from robogym.mujoco import constants as const from robogym.envs.rearrange.holdouts import STATE_DIR from robogym.viewer.robot_control_viewer import RobotControlViewer diff --git a/robogym/viewer/robot_control_viewer.py b/robogym/viewer/robot_control_viewer.py index 24859ee..221b8f9 100644 --- a/robogym/viewer/robot_control_viewer.py +++ b/robogym/viewer/robot_control_viewer.py @@ -3,7 +3,7 @@ import glfw import numpy as np from gym.spaces import MultiDiscrete -from mujoco_py import const +from robogym.mujoco import constants as const from robogym.robot.composite.controllers.ur_gripper_arm import ( Direction, diff --git a/setup.py b/setup.py index 07594d9..2a42292 100644 --- a/setup.py +++ b/setup.py @@ -1,5 +1,33 @@ #!/usr/bin/env python3 from setuptools import find_packages, setup +from distutils.extension import Extension +from os.path import dirname, realpath, join + + +SRC_DIR = dirname(realpath(__file__)) + + +class CallableList(object): + def __init__(self): + self.index = 0 + + def __len__(self): + return 1 + + def __getitem__(self, idx): + print("calling getitem!") + return None + + def __iter__(self): + return self + + def __next__(self): + if self.index == 1: + raise StopIteration() + if self.index == 0: + import mujoco + self.index += 1 + return join(mujoco.__path__[0], "include") def setup_robogym(): @@ -9,13 +37,13 @@ def setup_robogym(): packages=find_packages(), install_requires=[ # Fixed versions - "click==7.0", + "click", "collision==1.2.2", - "gym==0.15.3", + "gym", "kociemba==1.2.1", - "mujoco-py==2.0.2.13", + "mujoco>=2.2.0", "pycuber==0.2.2", - "matplotlib==3.1.2", + "matplotlib", "transforms3d==0.3.1", # Minimum versions "jsonnet>=0.14.0", @@ -27,6 +55,13 @@ def setup_robogym(): python_requires=">=3.7.4", description="OpenAI Robogym Robotics Environments", include_package_data=True, + ext_modules=[ + Extension('robogym.mujoco.callbacks', + sources=[join(SRC_DIR, "robogym", "mujoco", "callbacks.cpp")], + extra_compile_args=['-O3', '-std=c++17', '-Wno-#warnings', '-Wno-cpp', '-Wno-unused-function', '-Wno-deprecated-declarations'], + include_dirs=CallableList(), + language='c++') + ] ) From 7b5d5414a5edb30fdb97e3cf0150ffd6052cf26d Mon Sep 17 00:00:00 2001 From: Jonathan Raiman Date: Tue, 21 Jun 2022 10:28:09 -0700 Subject: [PATCH 2/8] working viewer --- robogym/mujoco/mujoco_xml.py | 235 +++++++++++++++++- robogym/mujoco/simulation_interface.py | 18 +- .../shadow_hand/mujoco/mujoco_shadow_hand.py | 2 +- robogym/viewer/env_viewer.py | 3 +- 4 files changed, 241 insertions(+), 17 deletions(-) diff --git a/robogym/mujoco/mujoco_xml.py b/robogym/mujoco/mujoco_xml.py index 481d657..617a53d 100644 --- a/robogym/mujoco/mujoco_xml.py +++ b/robogym/mujoco/mujoco_xml.py @@ -40,8 +40,65 @@ def __init__(self, model, data, nsubsteps=1): self._model = model self._data = data self._xml = None# model.get_xml() - self._nsubsteps = nsubsteps - print(f"nsubsteps = {nsubsteps}") + self.nsubsteps = nsubsteps + self._udd_callback = None + self._substep_callback = None + self.render_contexts = [] + self._render_context_offscreen = None + self._render_context_window = None + + def substep_callback(self): + if self._substep_callback is not None: + self._substep_callback(self._model.model(), self._data.data()) + + def reset(self): + mujoco.mj_resetData(self._model.model(), self._data.data()) + self.udd_state = None + self.step_udd() + + def add_render_context(self, render_context): + self.render_contexts.append(render_context) + if render_context.offscreen and self._render_context_offscreen is None: + self._render_context_offscreen = render_context + elif not render_context.offscreen and self._render_context_window is None: + self._render_context_window = render_context + + @property + def udd_callback(self): + return self._udd_callback + + @udd_callback.setter + def udd_callback(self, value): + self._udd_callback = value + self.udd_state = None + self.step_udd() + + def set_constants(self): + """ + Set constant fields of mjModel, corresponding to qpos0 configuration. + """ + mujoco.mj_setConst(self.model.model(), self.data.data()) + + def step_udd(self): + if self._udd_callback is None: + self.udd_state = {} + else: + schema_example = self.udd_state + self.udd_state = self._udd_callback(self) + # Check to make sure the udd_state has consistent keys and dimension across steps + if schema_example is not None: + keys = set(schema_example.keys()) | set(self.udd_state.keys()) + for key in keys: + assert key in schema_example, "Keys cannot be added to udd_state between steps." + assert key in self.udd_state, "Keys cannot be dropped from udd_state between steps." + if isinstance(schema_example[key], (int, float)): + assert isinstance(self.udd_state[key], (int, float)), \ + "Every value in udd_state must be either a number or a numpy array" + else: + assert isinstance(self.udd_state[key], np.ndarray), \ + "Every value in udd_state must be either a number or a numpy array" + assert self.udd_state[key].shape == schema_example[key].shape, \ + "Numpy array values in udd_state must keep the same dimension across steps." def get_xml(self): """ @@ -79,6 +136,176 @@ def _ensure_not_stale(self): "by the environment." ) + def forward(self): + return mujoco.mj_forward(self._model.model(), self._data.data()) + + def step(self, with_udd=True): + if with_udd: + self.step_udd() + for _ in range(self.nsubsteps): + self.substep_callback() + mujoco.mj_step(self.model.model(), self.data.data()) + + + +class DataProperty(object): + def __set_name__(self, owner, name): + self.name = name + + def __get__(self, obj, objtype=None): + if obj is None: + return self + return getattr(obj._data, self.name) + +class ModelProperty(object): + def __set_name__(self, owner, name): + self.name = name + + def __get__(self, obj, objtype=None): + if obj is None: + return self + return getattr(obj._model, self.name) + + +class MjModel(object): + def _extract_mj_names(self, name_adr, n, obj_type): + # objects don't need to be named in the XML, so name might be None + id2name = {i: None for i in range(n)} + name2id = {} + names = self._model.names + for i in range(n): + end = names.find(b"\x00", name_adr[i]) + name = names[name_adr[i]:end] + decoded_name = name.decode() + if decoded_name: + obj_id = mujoco.mj_name2id(self._model, obj_type, name) + assert 0 <= obj_id < n and id2name[obj_id] is None + name2id[decoded_name] = obj_id + id2name[obj_id] = decoded_name + + # sort names by increasing id to keep order deterministic + return tuple(id2name[id] for id in sorted(name2id.values())), name2id.__getitem__, name2id, id2name + + def __init__(self, model): + self._model = model + self.init_names() + + def model(self): + return self._model + + def init_names(self): + self.body_names, self.body_name2id, self._body_name2id, self._body_id2name = self._extract_mj_names(self._model.name_bodyadr, self._model.nbody, mujoco.mjtObj.mjOBJ_BODY) + self.joint_names, self.joint_name2id, self._joint_name2id, self._joint_id2name = self._extract_mj_names(self._model.name_jntadr, self._model.njnt, mujoco.mjtObj.mjOBJ_JOINT) + self.geom_names, self.geom_name2id, self._geom_name2id, self._geom_id2name = self._extract_mj_names(self._model.name_geomadr, self._model.ngeom, mujoco.mjtObj.mjOBJ_GEOM) + self.site_names, self.site_name2id, self._site_name2id, self._site_id2name = self._extract_mj_names(self._model.name_siteadr, self._model.nsite, mujoco.mjtObj.mjOBJ_SITE) + self.light_names, self.light_name2id, self._light_name2id, self._light_id2name = self._extract_mj_names(self._model.name_lightadr, self._model.nlight, mujoco.mjtObj.mjOBJ_LIGHT) + self.camera_names, self.camera_name2id, self._camera_name2id, self._camera_id2name = self._extract_mj_names(self._model.name_camadr, self._model.ncam, mujoco.mjtObj.mjOBJ_CAMERA) + self.actuator_names, self.actuator_name2id, self._actuator_name2id, self._actuator_id2name = self._extract_mj_names(self._model.name_actuatoradr, self._model.nu, mujoco.mjtObj.mjOBJ_ACTUATOR) + self.sensor_names, self.sensor_name2id, self._sensor_name2id, self._sensor_id2name = self._extract_mj_names(self._model.name_sensoradr, self._model.nsensor, mujoco.mjtObj.mjOBJ_SENSOR) + self.tendon_names, self.tendon_name2id, self._tendon_name2id, self._tendon_id2name = self._extract_mj_names(self._model.name_tendonadr, self._model.ntendon, mujoco.mjtObj.mjOBJ_TENDON) + self.mesh_names, self.mesh_name2id, self._mesh_name2id, self._mesh_id2name = self._extract_mj_names(self._model.name_meshadr, self._model.nmesh, mujoco.mjtObj.mjOBJ_MESH) + + def get_joint_qpos_addr(self, name): + ''' + Returns the qpos address for given joint. + Returns: + - address (int, tuple): returns int address if 1-dim joint, otherwise + returns the a (start, end) tuple for pos[start:end] access. + ''' + joint_id = self._joint_name2id[name] + joint_type = self._model.jnt_type[joint_id] + joint_addr = self._model.jnt_qposadr[joint_id] + if joint_type == mujoco.mjtJoint.mjJNT_FREE: + ndim = 7 + elif joint_type == mujoco.mjtJoint.mjJNT_BALL: + ndim = 4 + else: + assert joint_type in (mujoco.mjtJoint.mjJNT_HINGE, mujoco.mjtJoint.mjJNT_SLIDE) + ndim = 1 + + if ndim == 1: + return joint_addr + else: + return (joint_addr, joint_addr + ndim) + + actuator_forcerange = ModelProperty() + actuator_gainprm = ModelProperty() + actuator_ctrlrange = ModelProperty() + actuator_biasprm = ModelProperty() + body_pos = ModelProperty() + geom_size = ModelProperty() + body_inertia = ModelProperty() + geom_friction = ModelProperty() + site_xpos = ModelProperty() + site_pos = ModelProperty() + opt = ModelProperty() + tendon_range = ModelProperty() + nv = ModelProperty() + jnt_range = ModelProperty() + dof_jntid = ModelProperty() + dof_damping = ModelProperty() + body_mass = ModelProperty() + stat = ModelProperty() + ncam = ModelProperty() + geom_rgba = ModelProperty() + + def get_joint_qvel_addr(self, name): + ''' + Returns the qvel address for given joint. + Returns: + - address (int, tuple): returns int address if 1-dim joint, otherwise + returns the a (start, end) tuple for vel[start:end] access. + ''' + joint_id = self._joint_name2id[name] + joint_type = self._model.jnt_type[joint_id] + joint_addr = self._model.jnt_dofadr[joint_id] + if joint_type == mujoco.mjtJoint.mjJNT_FREE: + ndim = 6 + elif joint_type == mujoco.mjtJoint.mjJNT_BALL: + ndim = 3 + else: + assert joint_type in (mujoco.mjtJoint.mjJNT_HINGE, mujoco.mjtJoint.mjJNT_SLIDE) + ndim = 1 + + if ndim == 1: + return joint_addr + else: + return (joint_addr, joint_addr + ndim) + + +class MjData(object): + def __init__(self, data, model): + self._data = data + self._model = model + + def data(self): + return self._data + + def get_joint_qpos(self, name): + addr = self._model.get_joint_qpos_addr(name) + if isinstance(addr, (int, np.int32, np.int64)): + return self.qpos[addr] + else: + start_i, end_i = addr + return self.qpos[start_i:end_i] + + userdata = DataProperty() + actuator_force = DataProperty() + time = DataProperty() + qpos = DataProperty() + qvel = DataProperty() + ctrl = DataProperty() + site_xpos = DataProperty() + xfrc_applied = DataProperty() + ncon = DataProperty() + contact = DataProperty() + geom_xpos = DataProperty() + solver_iter = DataProperty() + + def get_site_xpos(self, name): + id = self._model.site_name2id(name) + return self._data.site_xpos[id] + class MujocoXML: """ @@ -245,8 +472,8 @@ def build(self, output_filename=None, meshdir=None, texturedir=None, **kwargs): with open(output_filename, "wt") as f: f.write(xml_string) - mj_model = mujoco.MjModel.from_xml_string(xml_string) - mj_data = mujoco.MjData(mj_model) + mj_model = MjModel(mujoco.MjModel.from_xml_string(xml_string)) + mj_data = MjData(mujoco.MjData(mj_model.model()), mj_model) return MjSim(mj_model, mj_data, **kwargs) ############################################################################################### diff --git a/robogym/mujoco/simulation_interface.py b/robogym/mujoco/simulation_interface.py index c2f8278..5bc74b3 100644 --- a/robogym/mujoco/simulation_interface.py +++ b/robogym/mujoco/simulation_interface.py @@ -87,13 +87,11 @@ def mujoco_viewer(self): """ Get a nicely-interactive version of the mujoco viewer """ - raise NotImplementedError() - # if self._mujoco_viewer is None: - # # Inline import since this is only relevant on platforms - # # which have GLFW support. - # from mujoco_py.mjviewer import MjViewer # noqa - - # self._mujoco_viewer = MjViewer(self.sim) + if self._mujoco_viewer is None: + # Inline import since this is only relevant on platforms + # which have GLFW support. + from .mjviewer import MjViewer # noqa + self._mujoco_viewer = MjViewer(self.sim) return self._mujoco_viewer @@ -102,13 +100,13 @@ def enable_pid(self): m = self.sim.model d = self.sim.data - if m.nuserdata < m.nu * NUM_USER_DATA_PER_ACT: + if m.model().nuserdata < m.model().nu * NUM_USER_DATA_PER_ACT: raise Exception('nuserdata is not set large enough to store PID internal states.') - if m.nuser_actuator < m.nu * NUM_ACTUATOR_DATA: + if m.model().nuser_actuator < m.model().nu * NUM_ACTUATOR_DATA: raise Exception('nuser_actuator is not set large enough to store controller types') - for i in range(m.nuserdata): + for i in range(m.model().nuserdata): d.userdata[i] = 0.0 mujoco.set_mjcb_act_gain(c_zero_gains) mujoco.set_mjcb_act_bias(c_custom_bias) diff --git a/robogym/robot/shadow_hand/mujoco/mujoco_shadow_hand.py b/robogym/robot/shadow_hand/mujoco/mujoco_shadow_hand.py index 9503a74..55722a1 100644 --- a/robogym/robot/shadow_hand/mujoco/mujoco_shadow_hand.py +++ b/robogym/robot/shadow_hand/mujoco/mujoco_shadow_hand.py @@ -86,7 +86,7 @@ def __init__( self.simulation.register_joint_group(self.joint_group, prefix=hand_prefix) self._parameter_manager = MuJoCoParameterManager(self.mj_sim) - assert self.mj_sim.model.nu == len( + assert self.mj_sim.model.model().nu == len( ACTUATORS ), "Action space must have compatible shape" diff --git a/robogym/viewer/env_viewer.py b/robogym/viewer/env_viewer.py index 7980a23..1eee860 100644 --- a/robogym/viewer/env_viewer.py +++ b/robogym/viewer/env_viewer.py @@ -6,8 +6,7 @@ import numpy as np from gym.spaces import Box, MultiDiscrete import robogym.mujoco.constants as const -from mujoco_py import MjViewer - +from robogym.mujoco.mjviewer import MjViewer from robogym.utils.misc import pretty From 9a6f99f9119d11786cf6907d5ed21de4cc4f244d Mon Sep 17 00:00:00 2001 From: Jonathan Raiman Date: Tue, 21 Jun 2022 10:29:56 -0700 Subject: [PATCH 3/8] add viewer files + demo script --- demo_viewer.py | 308 ++++++++++++++++++++++ robogym/mujoco/mjviewer.py | 409 ++++++++++++++++++++++++++++++ robogym/mujoco/mjviewercontext.py | 401 +++++++++++++++++++++++++++++ 3 files changed, 1118 insertions(+) create mode 100644 demo_viewer.py create mode 100644 robogym/mujoco/mjviewer.py create mode 100644 robogym/mujoco/mjviewercontext.py diff --git a/demo_viewer.py b/demo_viewer.py new file mode 100644 index 0000000..7d3a31c --- /dev/null +++ b/demo_viewer.py @@ -0,0 +1,308 @@ +from robogym.envs.dactyl.locked import make_env +import glfw +import mujoco +import numpy as np + +def init_window(max_width, max_height): + glfw.init() + window = glfw.create_window(width=max_width, height=max_height, + title='Demo', monitor=None, + share=None) + glfw.make_context_current(window) + return window + +# Fixed camera -1 is the free (unfixed) camera, and each fixed camera has +# a positive index in range (0, self._model.ncam). +_FREE_CAMERA_INDEX = -1 + +# Index used to distinguish when a camera isn't tracking any particular body. +_NO_BODY_TRACKED_INDEX = -1 + +# Index used to distinguish a non-existing or an invalid body. +_INVALID_BODY_INDEX = -1 + +# Zoom factor used when zooming in on the entire scene. +_FULL_SCENE_ZOOM_FACTOR = 1.5 + +class SceneCamera: + """A camera used to navigate around and render the scene.""" + + def __init__(self, + model, + data, + options, + settings=None, + zoom_factor=_FULL_SCENE_ZOOM_FACTOR): + """Instance initializer. + Args: + model: MjModel instance. + data: MjData instance. + options: RenderSettings instance. + settings: Optional, internal camera settings obtained from another + SceneCamera instance using 'settings' property. + zoom_factor: The initial zoom factor for zooming into the scene. + """ + # Design notes: + # We need to recreate the camera for each new model, because each model + # defines different fixed cameras and objects to track, and therefore + # severely the parameters of this class. + self._scene = wrapper.MjvScene(model) + self._data = data + self._model = model + self._options = options + + self._camera = wrapper.MjvCamera() + self._camera.trackbodyid = _NO_BODY_TRACKED_INDEX + self._camera.fixedcamid = _FREE_CAMERA_INDEX + self._camera.type_ = mujoco.mjtCamera.mjCAMERA_FREE + self._zoom_factor = zoom_factor + + if settings is not None: + self._settings = settings + self.settings = settings + else: + self._settings = self._camera + + def set_freelook_mode(self): + """Enables 6 degrees of freedom of movement for the camera.""" + self._camera.trackbodyid = _NO_BODY_TRACKED_INDEX + self._camera.fixedcamid = _FREE_CAMERA_INDEX + self._camera.type_ = mujoco.mjtCamera.mjCAMERA_FREE + + def set_tracking_mode(self, body_id): + """Latches the camera onto the specified body. + Leaves the user only 3 degrees of freedom to rotate the camera. + Args: + body_id: A positive integer, ID of the body to track. + """ + if body_id < 0: + return + self._camera.trackbodyid = body_id + self._camera.fixedcamid = _FREE_CAMERA_INDEX + self._camera.type_ = mujoco.mjtCamera.mjCAMERA_TRACKING + + def set_fixed_mode(self, fixed_camera_id): + """Fixes the camera in a pre-defined position, taking away all DOF. + Args: + fixed_camera_id: A positive integer, Id of a fixed camera defined in the + scene. + """ + if fixed_camera_id < 0: + return + self._camera.trackbodyid = _NO_BODY_TRACKED_INDEX + self._camera.fixedcamid = fixed_camera_id + self._camera.type_ = mujoco.mjtCamera.mjCAMERA_FIXED + + def look_at(self, position, distance): + """Positions the camera so that it's focused on the specified point.""" + self._camera.lookat[:] = position + self._camera.distance = distance + + def move(self, action, viewport_offset): + """Moves the camera around the scene.""" + # Not checking the validity of arguments on purpose. This method is designed + # to be called very often, so in order to avoid the overhead, all arguments + # are assumed to be valid. + mujoco.mjv_moveCamera(self._model.ptr, action, viewport_offset[0], + viewport_offset[1], self._scene.ptr, self._camera.ptr) + + def raycast(self, viewport, screen_pos): + """Shoots a ray from the specified viewport position into the scene.""" + if not self.is_initialized: + return -1, None + viewport_pos = viewport.screen_to_inverse_viewport(screen_pos) + grab_world_pos = np.empty(3, dtype=np.double) + selected_geom_id_arr = np.intc([-1]) + selected_skin_id_arr = np.intc([-1]) + selected_body_id = mujoco.mjv_select( + self._model.ptr, + self._data.ptr, + self._options.visualization.ptr, + viewport.aspect_ratio, + viewport_pos[0], + viewport_pos[1], + self._scene.ptr, + grab_world_pos, + selected_geom_id_arr, + selected_skin_id_arr, + ) + del selected_geom_id_arr, selected_skin_id_arr # Unused. + if selected_body_id < 0: + selected_body_id = _INVALID_BODY_INDEX + grab_world_pos = None + return selected_body_id, grab_world_pos + + def render(self, perturbation=None): + """Renders the scene form this camera's perspective. + Args: + perturbation: (Optional), instance of Perturbation. + Returns: + Rendered scene, instance of MjvScene. + """ + perturb_to_render = perturbation.ptr if perturbation else None + mujoco.mjv_updateScene(self._model.ptr, self._data.ptr, + self._options.visualization.ptr, perturb_to_render, + self._camera.ptr, mujoco.mjtCatBit.mjCAT_ALL, + self._scene.ptr) + return self._scene + + def zoom_to_scene(self): + """Zooms in on the entire scene.""" + self.look_at(self._model.stat.center[:], + self._zoom_factor * self._model.stat.extent) + + self.settings = self._settings + + @property + def transform(self): + """Returns a tuple with camera transform. + The transform comes in form: (3x3 rotation mtx, 3-component position). + """ + pos = np.zeros(3) + forward = np.zeros(3) + up = np.zeros(3) + for i in range(3): + forward[i] = self._scene.camera[0].forward[i] + up[i] = self._scene.camera[0].up[i] + pos[i] = (self._scene.camera[0].pos[i] + self._scene.camera[1].pos[i]) / 2 + right = np.cross(forward, up) + return np.array([right, up, forward]), pos + + @property + def settings(self): + """Returns internal camera settings.""" + return self._camera + + @settings.setter + def settings(self, value): + """Restores the camera settings.""" + self._camera.type_ = value.type_ + self._camera.fixedcamid = value.fixedcamid + self._camera.trackbodyid = value.trackbodyid + self._camera.lookat[:] = value.lookat[:] + self._camera.distance = value.distance + self._camera.azimuth = value.azimuth + self._camera.elevation = value.elevation + + @property + def name(self): + """Name of the active camera.""" + if self._camera.type_ == mujoco.mjtCamera.mjCAMERA_TRACKING: + body_name = self._model.id2name(self._camera.trackbodyid, 'body') + if body_name: + return 'Tracking body "%s"' % body_name + else: + return 'Tracking body id %d' % self._camera.trackbodyid + elif self._camera.type_ == mujoco.mjtCamera.mjCAMERA_FIXED: + camera_name = self._model.id2name(self._camera.fixedcamid, 'camera') + if camera_name: + return str(camera_name) + else: + return str(self._camera.fixedcamid) + else: + return 'Free' + + @property + def mode(self): + """Index of the mode the camera is currently in.""" + return self._camera.type_ + + @property + def is_initialized(self): + """Returns True if camera is properly initialized.""" + if not self._scene: + return False + frustum_near = self._scene.camera[0].frustum_near + frustum_far = self._scene.camera[0].frustum_far + return frustum_near > 0 and frustum_near < frustum_far + + +class MouseState(object): + def __init__(self): + self.pressed = False + + +def standalone_viewer(): + window = init_window(600, 450) + width, height = glfw.get_framebuffer_size(window) + scalex, _ = glfw.get_window_content_scale(window) + + + + if scalex != 1.0: + viewport = mujoco.MjrRect(-width, -height, width * 2, height * 2) + else: + viewport = mujoco.MjrRect(0, 0, width, height) + env = make_env() + obs = env.reset() + context = mujoco.MjrContext(env.unwrapped.mujoco_simulation.sim.model.model(), mujoco.mjtFontScale.mjFONTSCALE_100) + scene = mujoco.MjvScene(env.unwrapped.mujoco_simulation.sim.model.model(), 500) + camera = mujoco.MjvCamera() + mouse = MouseState() + def handle_scroll(window, x_offset, y_offset): + # scroll + mujoco.mjv_moveCamera(env.unwrapped.mujoco_simulation.sim.model.model(), mujoco.mjtMouse.mjMOUSE_ZOOM, 0, + 0.1 * y_offset, scene, camera) + + def button_change(window, button, state, mods): + if button == glfw.MOUSE_BUTTON_LEFT and state == glfw.PRESS: + mouse.pressed = True + mouse.startpos = mouse.last_position + elif button == glfw.MOUSE_BUTTON_LEFT and state != glfw.PRESS: + mouse.pressed = False + else: + pass + + def cursor_pos_change(window, posx, posy): + + if mouse.pressed: + dx = posx - mouse.last_position[0] + dy = posy - mouse.last_position[1] + mujoco.mjv_moveCamera(env.unwrapped.mujoco_simulation.sim.model.model(), mujoco.mjtMouse.mjMOUSE_MOVE_H, dx/width, + dy/height, scene, camera) + mouse.last_position = (posx, posy) + + glfw.set_scroll_callback(window, handle_scroll) + glfw.set_mouse_button_callback(window, button_change) + glfw.set_cursor_pos_callback(window, cursor_pos_change) + # import ipdb; ipdb.set_trace() + mujoco.mjv_updateScene( + env.unwrapped.mujoco_simulation.sim.model.model(), env.unwrapped.mujoco_simulation.sim.data.data(), + mujoco.MjvOption(), mujoco.MjvPerturb(), + camera, mujoco.mjtCatBit.mjCAT_ALL, scene) + i = 0 + while True and not glfw.window_should_close(window): + obs, reward, done, _ = env.step(env.action_space.nvec // 2) + mujoco.mjv_updateScene( + env.unwrapped.mujoco_simulation.sim.model.model(), env.unwrapped.mujoco_simulation.sim.data.data(), mujoco.MjvOption(), None, + camera, mujoco.mjtCatBit.mjCAT_ALL, scene) + mujoco.mjr_render(viewport, scene, context) + glfw.swap_buffers(window) + glfw.poll_events() + i += 1 + if i % 100 == 0: + env.reset() + + +def mujoco_viewer(): + env = make_env() + obs = env.reset() + i = 0 + while True: + obs, reward, done, _ = env.step(env.action_space.nvec // 2) + env.render() + i += 1 + if i % 100 == 0: + env.reset() + + +def main(): + standalone = False + if standalone: + standalone_viewer() + else: + mujoco_viewer() + + +if __name__ == "__main__": + main() diff --git a/robogym/mujoco/mjviewer.py b/robogym/mujoco/mjviewer.py new file mode 100644 index 0000000..e817a99 --- /dev/null +++ b/robogym/mujoco/mjviewer.py @@ -0,0 +1,409 @@ +import copy +import glfw +import imageio +import numpy as np +import time +import sys +from . import constants as const +from multiprocessing import Process, Queue +from threading import Lock +from .mjviewercontext import MjRenderContextWindow + + + +def rec_assign(node, assign): + # Assigns values to node recursively. + # This is neccessary to avoid overriding pointers in MuJoCo. + for field in dir(node): + if field.find("__") == -1 and field != 'uintptr': + val = getattr(node, field) + if isinstance(val, (int, bool, float, None.__class__, str)): + setattr(node, field, assign[field]) + elif isinstance(val, np.ndarray): + val[:] = assign[field][:] + elif not hasattr(val, "__call__"): + rec_assign(val, assign[field]) + + +def rec_copy(node): + # Recursively copies object to dictionary. + # Applying directly copy.deepcopy causes seg fault. + ret = {} + for field in dir(node): + if field.find("__") == -1: + val = getattr(node, field) + if isinstance(val, (int, bool, float, None.__class__, str)): + ret[field] = val + elif isinstance(val, np.ndarray): + ret[field] = copy.deepcopy(val) + elif not hasattr(val, "__call__"): + ret[field] = rec_copy(val) + return ret + + +class MjViewerBasic(MjRenderContextWindow): + """ + A simple display GUI showing the scene of an :class:`.MjSim` with a mouse-movable camera. + :class:`.MjViewer` extends this class to provide more sophisticated playback and interaction controls. + Parameters + ---------- + sim : :class:`.MjSim` + The simulator to display. + """ + + def __init__(self, sim): + super().__init__(sim) + + self._gui_lock = Lock() + self._button_left_pressed = False + self._button_right_pressed = False + self._last_mouse_x = 0 + self._last_mouse_y = 0 + + framebuffer_width, _ = glfw.get_framebuffer_size(self.window) + window_width, _ = glfw.get_window_size(self.window) + self._scale = framebuffer_width * 1.0 / window_width + + glfw.set_cursor_pos_callback(self.window, self._cursor_pos_callback) + glfw.set_mouse_button_callback( + self.window, self._mouse_button_callback) + glfw.set_scroll_callback(self.window, self._scroll_callback) + glfw.set_key_callback(self.window, self.key_callback) + + def render(self): + """ + Render the current simulation state to the screen or off-screen buffer. + Call this in your main loop. + """ + if self.window is None: + return + elif glfw.window_should_close(self.window): + glfw.terminate() + sys.exit(0) + + with self._gui_lock: + super().render() + + glfw.poll_events() + + def key_callback(self, window, key, scancode, action, mods): + if action == glfw.RELEASE and key == glfw.KEY_ESCAPE: + print("Pressed ESC") + print("Quitting.") + glfw.terminate() + sys.exit(0) + + def _cursor_pos_callback(self, window, xpos, ypos): + if not (self._button_left_pressed or self._button_right_pressed): + return + + # Determine whether to move, zoom or rotate view + mod_shift = ( + glfw.get_key(window, glfw.KEY_LEFT_SHIFT) == glfw.PRESS or + glfw.get_key(window, glfw.KEY_RIGHT_SHIFT) == glfw.PRESS) + if self._button_right_pressed: + action = const.MOUSE_MOVE_H if mod_shift else const.MOUSE_MOVE_V + elif self._button_left_pressed: + action = const.MOUSE_ROTATE_H if mod_shift else const.MOUSE_ROTATE_V + else: + action = const.MOUSE_ZOOM + + # Determine + dx = int(self._scale * xpos) - self._last_mouse_x + dy = int(self._scale * ypos) - self._last_mouse_y + width, height = glfw.get_framebuffer_size(window) + + with self._gui_lock: + self.move_camera(action, dx / height, dy / height) + + self._last_mouse_x = int(self._scale * xpos) + self._last_mouse_y = int(self._scale * ypos) + + def _mouse_button_callback(self, window, button, act, mods): + self._button_left_pressed = ( + glfw.get_mouse_button(window, glfw.MOUSE_BUTTON_LEFT) == glfw.PRESS) + self._button_right_pressed = ( + glfw.get_mouse_button(window, glfw.MOUSE_BUTTON_RIGHT) == glfw.PRESS) + + x, y = glfw.get_cursor_pos(window) + self._last_mouse_x = int(self._scale * x) + self._last_mouse_y = int(self._scale * y) + + def _scroll_callback(self, window, x_offset, y_offset): + with self._gui_lock: + self.move_camera(const.MOUSE_ZOOM, 0, -0.05 * y_offset) + + + +class MjViewer(MjViewerBasic): + """ + Extends :class:`.MjViewerBasic` to add video recording, interactive time and interaction controls. + The key bindings are as follows: + - TAB: Switch between MuJoCo cameras. + - H: Toggle hiding all GUI components. + - SPACE: Pause/unpause the simulation. + - RIGHT: Advance simulation by one step. + - V: Start/stop video recording. + - T: Capture screenshot. + - I: Drop into ``ipdb`` debugger. + - S/F: Decrease/Increase simulation playback speed. + - C: Toggle visualization of contact forces (off by default). + - D: Enable/disable frame skipping when rendering lags behind real time. + - R: Toggle transparency of geoms. + - M: Toggle display of mocap bodies. + - 0-4: Toggle display of geomgroups + Parameters + ---------- + sim : :class:`.MjSim` + The simulator to display. + """ + + def __init__(self, sim): + super().__init__(sim) + + self._ncam = sim.model.ncam + self._paused = False # is viewer paused. + # should we advance viewer just by one step. + self._advance_by_one_step = False + + # Vars for recording video + self._record_video = False + self._video_queue = Queue() + self._video_idx = 0 + self._video_path = "/tmp/video_%07d.mp4" + + # vars for capturing screen + self._image_idx = 0 + self._image_path = "/tmp/frame_%07d.png" + + # run_speed = x1, means running real time, x2 means fast-forward times + # two. + self._run_speed = 1.0 + self._loop_count = 0 + self._render_every_frame = False + + self._show_mocap = True # Show / hide mocap bodies. + self._transparent = False # Make everything transparent. + + # this variable is estamated as a running average. + self._time_per_render = 1 / 60.0 + self._hide_overlay = False # hide the entire overlay. + self._user_overlay = {} + + def render(self): + """ + Render the current simulation state to the screen or off-screen buffer. + Call this in your main loop. + """ + + def render_inner_loop(self): + render_start = time.time() + + self._overlay.clear() + if not self._hide_overlay: + for k, v in self._user_overlay.items(): + self._overlay[k] = copy.deepcopy(v) + self._create_full_overlay() + super().render() + if self._record_video: + frame = self._read_pixels_as_in_window() + self._video_queue.put(frame) + else: + self._time_per_render = 0.9 * self._time_per_render + \ + 0.1 * (time.time() - render_start) + + self._user_overlay = copy.deepcopy(self._overlay) + # Render the same frame if paused. + if self._paused: + while self._paused: + render_inner_loop(self) + if self._advance_by_one_step: + self._advance_by_one_step = False + break + else: + # inner_loop runs "_loop_count" times in expectation (where "_loop_count" is a float). + # Therefore, frames are displayed in the real-time. + self._loop_count += self.sim.model.opt.timestep * self.sim.nsubsteps / \ + (self._time_per_render * self._run_speed) + if self._render_every_frame: + self._loop_count = 1 + while self._loop_count > 0: + render_inner_loop(self) + self._loop_count -= 1 + # Markers and overlay are regenerated in every pass. + self._markers[:] = [] + self._overlay.clear() + + def _read_pixels_as_in_window(self, resolution=None): + # Reads pixels with markers and overlay from the same camera as screen. + if resolution is None: + resolution = glfw.get_framebuffer_size(self.sim._render_context_window.window) + + resolution = np.array(resolution) + resolution = resolution * min(1000 / np.min(resolution), 1) + resolution = resolution.astype(np.int32) + resolution -= resolution % 16 + if self.sim._render_context_offscreen is None: + self.sim.render(resolution[0], resolution[1]) + offscreen_ctx = self.sim._render_context_offscreen + window_ctx = self.sim._render_context_window + # Save markers and overlay from offscreen. + saved = [copy.deepcopy(offscreen_ctx._markers), + copy.deepcopy(offscreen_ctx._overlay), + rec_copy(offscreen_ctx.cam)] + # Copy markers and overlay from window. + offscreen_ctx._markers[:] = window_ctx._markers[:] + offscreen_ctx._overlay.clear() + offscreen_ctx._overlay.update(window_ctx._overlay) + rec_assign(offscreen_ctx.cam, rec_copy(window_ctx.cam)) + + img = self.sim.render(*resolution) + img = img[::-1, :, :] # Rendered images are upside-down. + # Restore markers and overlay to offscreen. + offscreen_ctx._markers[:] = saved[0][:] + offscreen_ctx._overlay.clear() + offscreen_ctx._overlay.update(saved[1]) + rec_assign(offscreen_ctx.cam, saved[2]) + return img + + def _create_full_overlay(self): + if self._render_every_frame: + self.add_overlay(const.GRID_TOPLEFT, "", "") + else: + self.add_overlay(const.GRID_TOPLEFT, "Run speed = %.3f x real time" % + self._run_speed, "[S]lower, [F]aster") + self.add_overlay( + const.GRID_TOPLEFT, "Ren[d]er every frame", "Off" if self._render_every_frame else "On") + self.add_overlay(const.GRID_TOPLEFT, "Switch camera (#cams = %d)" % (self._ncam + 1), + "[Tab] (camera ID = %d)" % self.cam.fixedcamid) + self.add_overlay(const.GRID_TOPLEFT, "[C]ontact forces", "Off" if self.vopt.flags[ + 10] == 1 else "On") + self.add_overlay( + const.GRID_TOPLEFT, "Referenc[e] frames", "Off" if self.vopt.frame == 1 else "On") + self.add_overlay(const.GRID_TOPLEFT, + "T[r]ansparent", "On" if self._transparent else "Off") + self.add_overlay( + const.GRID_TOPLEFT, "Display [M]ocap bodies", "On" if self._show_mocap else "Off") + if self._paused is not None: + if not self._paused: + self.add_overlay(const.GRID_TOPLEFT, "Stop", "[Space]") + else: + self.add_overlay(const.GRID_TOPLEFT, "Start", "[Space]") + self.add_overlay(const.GRID_TOPLEFT, + "Advance simulation by one step", "[right arrow]") + self.add_overlay(const.GRID_TOPLEFT, "[H]ide Menu", "") + if self._record_video: + ndots = int(7 * (time.time() % 1)) + dots = ("." * ndots) + (" " * (6 - ndots)) + self.add_overlay(const.GRID_TOPLEFT, + "Record [V]ideo (On) " + dots, "") + else: + self.add_overlay(const.GRID_TOPLEFT, "Record [V]ideo (Off) ", "") + if self._video_idx > 0: + fname = self._video_path % (self._video_idx - 1) + self.add_overlay(const.GRID_TOPLEFT, " saved as %s" % fname, "") + + self.add_overlay(const.GRID_TOPLEFT, "Cap[t]ure frame", "") + if self._image_idx > 0: + fname = self._image_path % (self._image_idx - 1) + self.add_overlay(const.GRID_TOPLEFT, " saved as %s" % fname, "") + self.add_overlay(const.GRID_TOPLEFT, "Start [i]pdb", "") + if self._record_video: + extra = " (while video is not recorded)" + else: + extra = "" + self.add_overlay(const.GRID_BOTTOMLEFT, "FPS", "%d%s" % + (1 / self._time_per_render, extra)) + self.add_overlay(const.GRID_BOTTOMLEFT, "Solver iterations", str( + self.sim.data.solver_iter + 1)) + step = round(self.sim.data.time / self.sim.model.opt.timestep) + self.add_overlay(const.GRID_BOTTOMRIGHT, "Step", str(step)) + self.add_overlay(const.GRID_BOTTOMRIGHT, "timestep", "%.5f" % self.sim.model.opt.timestep) + self.add_overlay(const.GRID_BOTTOMRIGHT, "n_substeps", str(self.sim.nsubsteps)) + self.add_overlay(const.GRID_TOPLEFT, "Toggle geomgroup visibility", "0-4") + + def key_callback(self, window, key, scancode, action, mods): + if action != glfw.RELEASE: + return + elif key == glfw.KEY_TAB: # Switches cameras. + self.cam.fixedcamid += 1 + self.cam.type = const.CAMERA_FIXED + if self.cam.fixedcamid >= self._ncam: + self.cam.fixedcamid = -1 + self.cam.type = const.CAMERA_FREE + elif key == glfw.KEY_H: # hides all overlay. + self._hide_overlay = not self._hide_overlay + elif key == glfw.KEY_SPACE and self._paused is not None: # stops simulation. + self._paused = not self._paused + # Advances simulation by one step. + elif key == glfw.KEY_RIGHT and self._paused is not None: + self._advance_by_one_step = True + self._paused = True + elif key == glfw.KEY_V or \ + (key == glfw.KEY_ESCAPE and self._record_video): # Records video. Trigers with V or if in progress by ESC. + self._record_video = not self._record_video + if self._record_video: + fps = (1 / self._time_per_render) + self._video_process = Process(target=save_video, + args=(self._video_queue, self._video_path % self._video_idx, fps)) + self._video_process.start() + if not self._record_video: + self._video_queue.put(None) + self._video_process.join() + self._video_idx += 1 + elif key == glfw.KEY_T: # capture screenshot + img = self._read_pixels_as_in_window() + imageio.imwrite(self._image_path % self._image_idx, img) + self._image_idx += 1 + elif key == glfw.KEY_I: # drops in debugger. + print('You can access the simulator by self.sim') + import ipdb + ipdb.set_trace() + elif key == glfw.KEY_S: # Slows down simulation. + self._run_speed /= 2.0 + elif key == glfw.KEY_F: # Speeds up simulation. + self._run_speed *= 2.0 + elif key == glfw.KEY_C: # Displays contact forces. + vopt = self.vopt + vopt.flags[10] = vopt.flags[11] = not vopt.flags[10] + elif key == glfw.KEY_D: # turn off / turn on rendering every frame. + self._render_every_frame = not self._render_every_frame + elif key == glfw.KEY_E: + vopt = self.vopt + vopt.frame = 1 - vopt.frame + elif key == glfw.KEY_R: # makes everything little bit transparent. + self._transparent = not self._transparent + if self._transparent: + self.sim.model.geom_rgba[:, 3] /= 5.0 + else: + self.sim.model.geom_rgba[:, 3] *= 5.0 + elif key == glfw.KEY_M: # Shows / hides mocap bodies + self._show_mocap = not self._show_mocap + for body_idx1, val in enumerate(self.sim.model.body_mocapid): + if val != -1: + for geom_idx, body_idx2 in enumerate(self.sim.model.geom_bodyid): + if body_idx1 == body_idx2: + if not self._show_mocap: + # Store transparency for later to show it. + self.sim.extras[ + geom_idx] = self.sim.model.geom_rgba[geom_idx, 3] + self.sim.model.geom_rgba[geom_idx, 3] = 0 + else: + self.sim.model.geom_rgba[ + geom_idx, 3] = self.sim.extras[geom_idx] + elif key in (glfw.KEY_0, glfw.KEY_1, glfw.KEY_2, glfw.KEY_3, glfw.KEY_4): + self.vopt.geomgroup[key - glfw.KEY_0] ^= 1 + super().key_callback(window, key, scancode, action, mods) + +# Separate Process to save video. This way visualization is +# less slowed down. + + +def save_video(queue, filename, fps): + writer = imageio.get_writer(filename, fps=fps) + while True: + frame = queue.get() + if frame is None: + break + writer.append_data(frame) + writer.close() diff --git a/robogym/mujoco/mjviewercontext.py b/robogym/mujoco/mjviewercontext.py new file mode 100644 index 0000000..aca4ffc --- /dev/null +++ b/robogym/mujoco/mjviewercontext.py @@ -0,0 +1,401 @@ +import os +import numpy as np +import glfw +import sys +from abc import ABCMeta, abstractmethod +from . import constants as const +import mujoco + + +class OpenGLContext(metaclass=ABCMeta): + + @abstractmethod + def make_context_current(self): + raise NotImplementedError() + + @abstractmethod + def set_buffer_size(self, width, height): + raise NotImplementedError() + + +class GlfwError(RuntimeError): + pass + + +class GlfwContext(OpenGLContext): + + _INIT_WIDTH = 1000 + _INIT_HEIGHT = 1000 + _GLFW_IS_INITIALIZED = False + + def __init__(self, offscreen=False, quiet=False): + GlfwContext._init_glfw() + + self._width = self._INIT_WIDTH + self._height = self._INIT_HEIGHT + self.window = self._create_window(offscreen, quiet=quiet) + self._set_window_size(self._width, self._height) + + @staticmethod + def _init_glfw(): + if GlfwContext._GLFW_IS_INITIALIZED: + return + + if 'glfw' not in globals(): + raise GlfwError("GLFW not installed") + + glfw.set_error_callback(GlfwContext._glfw_error_callback) + + # HAX: sometimes first init() fails, while second works fine. + glfw.init() + if not glfw.init(): + raise GlfwError("Failed to initialize GLFW") + + GlfwContext._GLFW_IS_INITIALIZED = True + + def make_context_current(self): + glfw.make_context_current(self.window) + + def set_buffer_size(self, width, height): + self._set_window_size(width, height) + self._width = width + self._height = height + + def _create_window(self, offscreen, quiet=True): + if offscreen: + if not quiet: + print("Creating offscreen glfw") + glfw.window_hint(glfw.VISIBLE, 0) + glfw.window_hint(glfw.DOUBLEBUFFER, 0) + init_width, init_height = self._INIT_WIDTH, self._INIT_HEIGHT + else: + if not quiet: + print("Creating window glfw") + glfw.window_hint(glfw.SAMPLES, 4) + glfw.window_hint(glfw.VISIBLE, 1) + glfw.window_hint(glfw.DOUBLEBUFFER, 1) + resolution, _, refresh_rate = glfw.get_video_mode( + glfw.get_primary_monitor()) + init_width, init_height = resolution + + self._width = init_width + self._height = init_height + window = glfw.create_window( + self._width, self._height, "robogym", None, None) + + if not window: + raise GlfwError("Failed to create GLFW window") + + return window + + def get_buffer_size(self): + return glfw.get_framebuffer_size(self.window) + + def _set_window_size(self, target_width, target_height): + self.make_context_current() + if target_width != self._width or target_height != self._height: + self._width = target_width + self._height = target_height + glfw.set_window_size(self.window, target_width, target_height) + + # HAX: When running on a Mac with retina screen, the size + # sometimes doubles + width, height = glfw.get_framebuffer_size(self.window) + if target_width != width and "darwin" in sys.platform.lower(): + glfw.set_window_size(self.window, target_width // 2, target_height // 2) + + @staticmethod + def _glfw_error_callback(error_code, description): + print("GLFW error (code %d): %s", error_code, description) + + +# TODO +# class OffscreenOpenGLContext(): +# def __init__(self, device_id): +# self.device_id = device_id +# res = initOpenGL(device_id) +# if res != 1: +# raise RuntimeError("Failed to initialize OpenGL") + +# def close(self): +# # TODO: properly close OpenGL in our contexts +# closeOpenGL() + +# def make_context_current(self): +# makeOpenGLContextCurrent(self.device_id) + +# def set_buffer_size(self, int width, int height): +# res = setOpenGLBufferSize(self.device_id, width, height) +# if res != 1: +# raise RuntimeError("Failed to set buffer size") + + +class MjRenderContext(object): + """ + Class that encapsulates rendering functionality for a + MuJoCo simulation. + """ + + def __init__(self, sim, offscreen=True, device_id=-1, opengl_backend=None, quiet=False): + self.sim = sim + self._setup_opengl_context(offscreen, device_id, opengl_backend, quiet=quiet) + self.offscreen = offscreen + + # Ensure the model data has been updated so that there + # is something to render + sim.forward() + sim.add_render_context(self) + self._model_ptr = sim.model.model() + self._data_ptr = sim.data.data() + maxgeom = 1000 + self.scn = mujoco.MjvScene(self._model_ptr, maxgeom) + self.cam = mujoco.MjvCamera() + mujoco.mjv_defaultCamera(self.cam) + self.vopt = mujoco.MjvOption() + self.con = None + + self.pert = mujoco.MjvPerturb() + self.pert.active = 0 + self.pert.select = 0 + self.pert.skinselect = -1 + + self._markers = [] + self._overlay = {} + + self._init_camera(sim) + self._set_mujoco_buffers() + + def update_sim(self, new_sim): + if new_sim == self.sim: + return + self._model_ptr = new_sim.model.model() + self._data_ptr = new_sim.data.data() + self._set_mujoco_buffers() + for render_context in self.sim.render_contexts: + new_sim.add_render_context(render_context) + self.sim = new_sim + + def _set_mujoco_buffers(self): + self.con = mujoco.MjrContext(self._model_ptr, mujoco.mjtFontScale.mjFONTSCALE_150) + if self.offscreen: + mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_OFFSCREEN, self.con) + if self.con.currentBuffer != mujoco.mjtFramebuffer.mjFB_OFFSCREEN: + raise RuntimeError("Offscreen rendering not supported") + else: + mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_WINDOW, self.con); + if self.con.currentBuffer != mujoco.mjtFramebuffer.mjFB_WINDOW: + raise RuntimeError("Window rendering not supported") + + def _setup_opengl_context(self, offscreen, device_id, opengl_backend, quiet=False): + if opengl_backend is None and (not offscreen or sys.platform == 'darwin'): + # default to glfw for onscreen viewing or mac (both offscreen/onscreen) + opengl_backend = 'glfw' + + if opengl_backend == 'glfw': + self.opengl_context = GlfwContext(offscreen=offscreen, quiet=quiet) + else: + raise NotImplementedError() + if device_id < 0: + if "GPUS" in os.environ: + device_id = os.environ["GPUS"] + else: + device_id = os.getenv('CUDA_VISIBLE_DEVICES', '') + if len(device_id) > 0: + device_id = int(device_id.split(',')[0]) + else: + # Sometimes env variable is an empty string. + device_id = 0 + self.opengl_context = OffscreenOpenGLContext(device_id) + + def _init_camera(self, sim): + # Make the free camera look at the scene + self.cam.type = const.CAMERA_FREE + self.cam.fixedcamid = -1 + for i in range(3): + self.cam.lookat[i] = np.median(sim.data.geom_xpos[:, i]) + self.cam.distance = sim.model.stat.extent + + def update_offscreen_size(self, width, height): + if width != self.con.offWidth or height != self.con.offHeight: + self._model_ptr.vis.global_.offwidth = width + self._model_ptr.vis.global_.offheight = height + # mujoco.mjr_freeContext(self._con) + self._set_mujoco_buffers() + + def render(self, width, height, camera_id=None, segmentation=False): + rect = mujoco.MjrRect(0, 0, width, height) + + if self.sim.render_callback is not None: + self.sim.render_callback(self.sim, self) + + # Sometimes buffers are too small. + if width > self.con.offWidth or height > self.con.offHeight: + new_width = max(width, self._model_ptr.vis.global_.offwidth) + new_height = max(height, self._model_ptr.vis.global_.offheight) + self.update_offscreen_size(new_width, new_height) + + if camera_id is not None: + if camera_id == -1: + self.cam.type = const.CAMERA_FREE + else: + self.cam.type = const.CAMERA_FIXED + self.cam.fixedcamid = camera_id + + # This doesn't really do anything else rather than checking for the size of buffer + # need to investigate further whi is that a no-op + # self.opengl_context.set_buffer_size(width, height) + + mujoco.mjv_updateScene( + self._model_ptr, self._data_ptr, + self.vopt, self.pert, + self.cam, mujoco.mjtCatBit.mjCAT_ALL, self.scn) + + if segmentation: + self.scn.flags[const.RND_SEGMENT] = 1 + self.scn.flags[const.RND_IDCOLOR] = 1 + + for marker_params in self._markers: + self._add_marker_to_scene(marker_params) + + mujoco.mjr_render(rect, self.scn, self.con) + for gridpos, (text1, text2) in self._overlay.items(): + mujoco.mjr_overlay(const.FONTSCALE_150, gridpos, rect, text1.encode(), text2.encode(), self.con) + + if segmentation: + self.scn.flags[const.RND_SEGMENT] = 0 + self.scn.flags[const.RND_IDCOLOR] = 0 + + def read_pixels(self, width, height, depth=True, segmentation=False): + rect = mujoco.MjrRect(0, 0, width, height) + + rgb_arr = np.zeros(3 * rect.width * rect.height, dtype=np.uint8) + depth_arr = np.zeros(rect.width * rect.height, dtype=np.float32) + + mujoco.mjr_readPixels(rgb_arr, depth_arr, rect, self.con) + rgb_img = rgb_arr.reshape(rect.height, rect.width, 3) + ret_img = rgb_img + if segmentation: + seg_img = (rgb_img[:, :, 0] + rgb_img[:, :, 1] * (2**8) + rgb_img[:, :, 2] * (2 ** 16)) + seg_img[seg_img >= (self.scn.ngeom + 1)] = 0 + seg_ids = np.full((self.scn.ngeom + 1, 2), fill_value=-1, dtype=np.int32) + + for i in range(self.scn.ngeom): + geom = self.scn.geoms[i] + if geom.segid != -1: + seg_ids[geom.segid + 1, 0] = geom.objtype + seg_ids[geom.segid + 1, 1] = geom.objid + ret_img = seg_ids[seg_img] + + if depth: + depth_img = depth_arr.reshape(rect.height, rect.width) + return (ret_img, depth_img) + else: + return ret_img + + def read_pixels_depth(self, buffer): + ''' Read depth pixels into a preallocated buffer ''' + rect = mujoco.MjrRect(0, 0, buffer.shape[1], buffer.shape[0]) + mujoco.mjr_readPixels(0, buffer, rect, self.con) + + def upload_texture(self, tex_id): + """ Uploads given texture to the GPU. """ + self.opengl_context.make_context_current() + mujoco.mjr_uploadTexture(self._model_ptr, self.con, tex_id) + + def draw_pixels(self, image, left, bottom): + """Draw an image into the OpenGL buffer.""" + assert isinstance(image, np.ndarray), f"expected image to be of type np.ndarray but got {type(image)}." + assert image.dtype == np.uint8, f"expected image to be of dtype np.uint8 but got {image.dtype}." + assert image.ndim == 3, f"expected image to have ndim=3 but got {image.ndim}." + viewport = mujoco.MjrRect(0, 0, image.shape[1], image.shape[0]) + mujoco.mjr_drawPixels(image, 0, viewport, self.con) + + def move_camera(self, action, reldx, reldy): + """ Moves the camera based on mouse movements. Action is one of mjMOUSE_*. """ + mujoco.mjv_moveCamera(self._model_ptr, action, reldx, reldy, self.scn, self.cam) + + def add_overlay(self, gridpos, text1, text2): + """ Overlays text on the scene. """ + if gridpos not in self._overlay: + self._overlay[gridpos] = ["", ""] + self._overlay[gridpos][0] += text1 + "\n" + self._overlay[gridpos][1] += text2 + "\n" + + def add_marker(self, **marker_params): + self._markers.append(marker_params) + + def _add_marker_to_scene(self, marker_params): + """ Adds marker to scene, and returns the corresponding object. """ + if self.scn.ngeom >= self.scn.maxgeom: + raise RuntimeError('Ran out of geoms. maxgeom: %d' % self.scn.maxgeom) + + # cdef mjvGeom *g = self.scn.geoms + self.scn.ngeom + g = mujoco.MjvGeom() + + # default values. + g.dataid = -1 + g.objtype = const.OBJ_UNKNOWN + g.objid = -1 + g.category = const.CAT_DECOR + g.texid = -1 + g.texuniform = 0 + g.texrepeat[0] = 1 + g.texrepeat[1] = 1 + g.emission = 0 + g.specular = 0.5 + g.shininess = 0.5 + g.reflectance = 0 + g.type = const.GEOM_BOX + g.size[:] = np.ones(3) * 0.1 + g.mat[:] = np.eye(3).flatten() + g.rgba[:] = np.ones(4) + + for key, value in marker_params.items(): + if isinstance(value, (int, float)): + setattr(g, key, value) + elif isinstance(value, (tuple, list, np.ndarray)): + attr = getattr(g, key) + attr[:] = np.asarray(value).reshape(attr.shape) + elif isinstance(value, str): + assert key == "label", "Only label is a string in mjvGeom." + if value == None: + g.label = 0 + else: + g.label = value.encode() + elif hasattr(g, key): + raise ValueError("mjvGeom has attr {} but type {} is invalid".format(key, type(value))) + else: + raise ValueError("mjvGeom doesn't have field %s" % key) + + self.scn.ngeom += 1 + # def __dealloc__(self): + # mjr_freeContext(&self._con) + # mjv_freeScene(&self._scn) + + +class MjRenderContextOffscreen(MjRenderContext): + def __cinit__(self, sim, device_id): + super().__init__(sim, offscreen=True, device_id=device_id) + + +class MjRenderContextWindow(MjRenderContext): + def __init__(self, sim): + super().__init__(sim, offscreen=False) + self.render_swap_callback = None + + assert isinstance(self.opengl_context, GlfwContext), ( + "Only GlfwContext supported for windowed rendering") + + @property + def window(self): + return self.opengl_context.window + + def render(self): + if self.window is None or glfw.window_should_close(self.window): + return + + glfw.make_context_current(self.window) + super().render(*glfw.get_framebuffer_size(self.window)) + if self.render_swap_callback is not None: + self.render_swap_callback() + glfw.swap_buffers(self.window) From 8877355f9f9c5e9d55288fee61f6574e5af4f722 Mon Sep 17 00:00:00 2001 From: Jonathan Raiman Date: Wed, 22 Jun 2022 17:28:34 -0700 Subject: [PATCH 4/8] add support for reach env with cont reward --- raptor_configs/reach.jsonnet | 11 +++++++++ raptor_robogym/__init__.py | 46 ++++++++++++++++++++++++++++++++++++ robogym/envs/dactyl/reach.py | 6 ++--- robogym/mujoco/mujoco_xml.py | 1 + 4 files changed, 61 insertions(+), 3 deletions(-) create mode 100644 raptor_configs/reach.jsonnet create mode 100644 raptor_robogym/__init__.py diff --git a/raptor_configs/reach.jsonnet b/raptor_configs/reach.jsonnet new file mode 100644 index 0000000..4aa55fe --- /dev/null +++ b/raptor_configs/reach.jsonnet @@ -0,0 +1,11 @@ +local base = import 'raptor/experiments/base.jsonnet'; +base.rl_training( + env={"function": "raptor_robogym:wrap_env", + "args": {"env": "robogym.envs.dactyl.reach:make_env"}}, + model="rl.learn.policy:make_mlp_policy", + algorithm="rl.learn.actor_critic:PPO", + use_distributed_communication=false, + buffer_size=1024, + batch_size=256, + workers=8 +) \ No newline at end of file diff --git a/raptor_robogym/__init__.py b/raptor_robogym/__init__.py new file mode 100644 index 0000000..8990e59 --- /dev/null +++ b/raptor_robogym/__init__.py @@ -0,0 +1,46 @@ +import gym +import numpy as np +from rl.utils.make_callable import make_callable + +class NormReward(gym.Wrapper): + def step(self, action): + obs, reward, done, info = self.env.step(action) + return obs, sum(reward), done, info + +class DictToArray(gym.Wrapper): + def __init__(self, env): + super(DictToArray, self).__init__(env) + assert isinstance(self.env.observation_space, gym.spaces.Dict), \ + f"expected wrapped env's observation space to be a gym.spaces.Dict but got {self.env.observation_space}." + ignored_keys = {"performed_action", "action_history", "action_delay", "randomized_phasespace", "randomized_phasespace_fingers"} + self._space_names = [n for n in self.env.observation_space.spaces.keys() if n not in ignored_keys] + self._space_shapes = [np.prod(s.shape) for n, s in self.env.observation_space.spaces.items() if n not in ignored_keys] + self.observation_space = gym.spaces.Box(low=-float("inf"), high=float("inf"), shape=(sum(self._space_shapes),)) + + def reset(self): + return self._concate_obs(self.env.reset()) + + def step(self, action): + obs, reward, done, info = self.env.step(action) + return self._concate_obs(obs), reward, done, info + + def _concate_obs(self, observation): + out = np.zeros(self.observation_space.shape[0], dtype=np.float32) + so_far = 0 + for key, size in zip(self._space_names, self._space_shapes): + if key in observation: + obs = observation[key].reshape(-1) + assert len(obs) == size, \ + f"expected observation named {key} to have size {size} but found {len(obs)} instead (observation[key].shape = {observation[key].shape} -- {self.env.observation_space[key]})." + out[so_far:so_far + size] = obs + so_far += size + else: + so_far += size + print(f"wtf! {key}") + return out + + + +def wrap_env(env): + env = env if isinstance(env, gym.Env) else make_callable(env)() + return NormReward(DictToArray(env)) diff --git a/robogym/envs/dactyl/reach.py b/robogym/envs/dactyl/reach.py index 3b38d6b..94db327 100644 --- a/robogym/envs/dactyl/reach.py +++ b/robogym/envs/dactyl/reach.py @@ -251,9 +251,9 @@ def apply_wrappers(self, **wrapper_params): env = util.ClipObservationWrapper(env) env = util.ClipRewardWrapper(env) env = util.PreviousActionObservationWrapper(env) - env = util.DiscretizeActionWrapper( - env, n_action_bins=self.constants.n_action_bins - ) + # env = util.DiscretizeActionWrapper( + # env, n_action_bins=self.constants.n_action_bins + # ) # Note: Recording wrapper is removed here to favor simplicity. return env diff --git a/robogym/mujoco/mujoco_xml.py b/robogym/mujoco/mujoco_xml.py index 617a53d..1ab2f31 100644 --- a/robogym/mujoco/mujoco_xml.py +++ b/robogym/mujoco/mujoco_xml.py @@ -248,6 +248,7 @@ def get_joint_qpos_addr(self, name): stat = ModelProperty() ncam = ModelProperty() geom_rgba = ModelProperty() + geom_margin = ModelProperty() def get_joint_qvel_addr(self, name): ''' From 37871dc659e0a5135f5101519a07098e16877188 Mon Sep 17 00:00:00 2001 From: Jonathan Raiman Date: Wed, 22 Jun 2022 20:22:32 -0700 Subject: [PATCH 5/8] update conf --- raptor_configs/reach.jsonnet | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/raptor_configs/reach.jsonnet b/raptor_configs/reach.jsonnet index 4aa55fe..1a2cd3c 100644 --- a/raptor_configs/reach.jsonnet +++ b/raptor_configs/reach.jsonnet @@ -2,10 +2,13 @@ local base = import 'raptor/experiments/base.jsonnet'; base.rl_training( env={"function": "raptor_robogym:wrap_env", "args": {"env": "robogym.envs.dactyl.reach:make_env"}}, - model="rl.learn.policy:make_mlp_policy", - algorithm="rl.learn.actor_critic:PPO", + model={"function": "rl.learn.policy:make_mlp_policy", "args": {"normalize": false}}, + algorithm={"function": "rl.learn.actor_critic:PPO", "args": {"advantage_normalize": true}}, + optimizer_fn=base.optimizer_fn(optim_stepsize=3e-5), use_distributed_communication=false, - buffer_size=1024, + optimizer_clip_norm=1.0, + buffer_size=256, + max_staleness=0, batch_size=256, workers=8 ) \ No newline at end of file From 4c594103c504195d18c0cd9a515538a3a550885d Mon Sep 17 00:00:00 2001 From: Jonathan Raiman Date: Wed, 22 Jun 2022 20:34:12 -0700 Subject: [PATCH 6/8] simplify obs --- demo_viewer.py | 6 ++++-- raptor_configs/reach.jsonnet | 2 +- raptor_robogym/__init__.py | 5 ++--- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/demo_viewer.py b/demo_viewer.py index 7d3a31c..71c9748 100644 --- a/demo_viewer.py +++ b/demo_viewer.py @@ -1,4 +1,5 @@ -from robogym.envs.dactyl.locked import make_env +# from robogym.envs.dactyl.locked import make_env +from robogym.envs.dactyl.reach import make_env import glfw import mujoco import numpy as np @@ -272,6 +273,7 @@ def cursor_pos_change(window, posx, posy): camera, mujoco.mjtCatBit.mjCAT_ALL, scene) i = 0 while True and not glfw.window_should_close(window): + print(env.action_space.sample()) obs, reward, done, _ = env.step(env.action_space.nvec // 2) mujoco.mjv_updateScene( env.unwrapped.mujoco_simulation.sim.model.model(), env.unwrapped.mujoco_simulation.sim.data.data(), mujoco.MjvOption(), None, @@ -289,7 +291,7 @@ def mujoco_viewer(): obs = env.reset() i = 0 while True: - obs, reward, done, _ = env.step(env.action_space.nvec // 2) + obs, reward, done, _ = env.step(env.action_space.sample() * 0) env.render() i += 1 if i % 100 == 0: diff --git a/raptor_configs/reach.jsonnet b/raptor_configs/reach.jsonnet index 1a2cd3c..4af0096 100644 --- a/raptor_configs/reach.jsonnet +++ b/raptor_configs/reach.jsonnet @@ -1,7 +1,7 @@ local base = import 'raptor/experiments/base.jsonnet'; base.rl_training( env={"function": "raptor_robogym:wrap_env", - "args": {"env": "robogym.envs.dactyl.reach:make_env"}}, + "args": {"env": {"function": "robogym.envs.dactyl.reach:make_env", "args": {"constants": {"randomize": false}}},}}, model={"function": "rl.learn.policy:make_mlp_policy", "args": {"normalize": false}}, algorithm={"function": "rl.learn.actor_critic:PPO", "args": {"advantage_normalize": true}}, optimizer_fn=base.optimizer_fn(optim_stepsize=3e-5), diff --git a/raptor_robogym/__init__.py b/raptor_robogym/__init__.py index 8990e59..c808ab0 100644 --- a/raptor_robogym/__init__.py +++ b/raptor_robogym/__init__.py @@ -12,9 +12,8 @@ def __init__(self, env): super(DictToArray, self).__init__(env) assert isinstance(self.env.observation_space, gym.spaces.Dict), \ f"expected wrapped env's observation space to be a gym.spaces.Dict but got {self.env.observation_space}." - ignored_keys = {"performed_action", "action_history", "action_delay", "randomized_phasespace", "randomized_phasespace_fingers"} - self._space_names = [n for n in self.env.observation_space.spaces.keys() if n not in ignored_keys] - self._space_shapes = [np.prod(s.shape) for n, s in self.env.observation_space.spaces.items() if n not in ignored_keys] + self._space_names = ["qpos", "qvel", "fingertip_pos", "goal_fingertip_pos"] + self._space_shapes = [np.prod(self.env.observation_space.spaces[n].shape) for n in self._space_names] self.observation_space = gym.spaces.Box(low=-float("inf"), high=float("inf"), shape=(sum(self._space_shapes),)) def reset(self): From b338e3a0100b776e39b460fbeb217c8d6a3ba28d Mon Sep 17 00:00:00 2001 From: Jonathan Raiman Date: Thu, 23 Jun 2022 15:06:06 -0700 Subject: [PATCH 7/8] fix installf for windows --- robogym/mujoco/callbacks.cpp | 2 +- setup.py | 7 ++++++- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/robogym/mujoco/callbacks.cpp b/robogym/mujoco/callbacks.cpp index a034243..5c9bb9c 100644 --- a/robogym/mujoco/callbacks.cpp +++ b/robogym/mujoco/callbacks.cpp @@ -240,7 +240,7 @@ extern "C" mjtNum c_pi_cascade_bias(const mjModel* m, const mjData* d, int id) // Gravity compensation f += d->qfrc_bias[id]; - if (effort_limit_low != 0.0 or effort_limit_high != 0.0) { + if (effort_limit_low != 0.0 || effort_limit_high != 0.0) { f = fmax(effort_limit_low, fmin(effort_limit_high, f)); } return f; diff --git a/setup.py b/setup.py index 2a42292..46515e0 100644 --- a/setup.py +++ b/setup.py @@ -2,6 +2,7 @@ from setuptools import find_packages, setup from distutils.extension import Extension from os.path import dirname, realpath, join +import platform SRC_DIR = dirname(realpath(__file__)) @@ -31,6 +32,10 @@ def __next__(self): def setup_robogym(): + if platform.system() == "Windows": + extra_compile_args = ['/O2', '-std:c++17'] + else: + extra_compile_args = ['-O3', '-std=c++17', '-Wno-#warnings', '-Wno-cpp', '-Wno-unused-function', '-Wno-deprecated-declarations'] setup( name="robogym", version=open("ROBOGYM_VERSION").read(), @@ -58,7 +63,7 @@ def setup_robogym(): ext_modules=[ Extension('robogym.mujoco.callbacks', sources=[join(SRC_DIR, "robogym", "mujoco", "callbacks.cpp")], - extra_compile_args=['-O3', '-std=c++17', '-Wno-#warnings', '-Wno-cpp', '-Wno-unused-function', '-Wno-deprecated-declarations'], + extra_compile_args=extra_compile_args, include_dirs=CallableList(), language='c++') ] From f92f53021f0b3cb82eb32dba54a09ab2222d0382 Mon Sep 17 00:00:00 2001 From: Jonathan Raiman Date: Mon, 27 Jun 2022 12:22:42 -0700 Subject: [PATCH 8/8] lazy extension --- setup.py | 39 ++++++++++++++------------------------- 1 file changed, 14 insertions(+), 25 deletions(-) diff --git a/setup.py b/setup.py index 46515e0..e7d0610 100644 --- a/setup.py +++ b/setup.py @@ -8,27 +8,17 @@ SRC_DIR = dirname(realpath(__file__)) -class CallableList(object): - def __init__(self): - self.index = 0 +class LazyExtension(Extension): + """enable lazy binding of the include directory once we know that mujoco has been installed + and we know the right path.""" + @property + def include_dirs(self): + import mujoco + return [join(mujoco.__path__[0], "include")] - def __len__(self): - return 1 - - def __getitem__(self, idx): - print("calling getitem!") - return None - - def __iter__(self): - return self - - def __next__(self): - if self.index == 1: - raise StopIteration() - if self.index == 0: - import mujoco - self.index += 1 - return join(mujoco.__path__[0], "include") + @include_dirs.setter + def include_dirs(self, value): + pass def setup_robogym(): @@ -61,11 +51,10 @@ def setup_robogym(): description="OpenAI Robogym Robotics Environments", include_package_data=True, ext_modules=[ - Extension('robogym.mujoco.callbacks', - sources=[join(SRC_DIR, "robogym", "mujoco", "callbacks.cpp")], - extra_compile_args=extra_compile_args, - include_dirs=CallableList(), - language='c++') + LazyExtension('robogym.mujoco.callbacks', + sources=[join(SRC_DIR, "robogym", "mujoco", "callbacks.cpp")], + extra_compile_args=extra_compile_args, + language='c++') ] )