Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Best practices for fast rendering with MJX #1604

Open
c0g opened this issue Apr 18, 2024 · 3 comments
Open

Best practices for fast rendering with MJX #1604

c0g opened this issue Apr 18, 2024 · 3 comments
Labels
question Request for help or information

Comments

@c0g
Copy link

c0g commented Apr 18, 2024

Hi,

I'm a engineer and I'm trying to use MuJoCo for fun.

I'm looking for some help with rendering simulations in MJX. My model is extremely simple so this might be part of the issue, e.g. I'm just not doing enough compute (model at the bottom of this post).

I've verified that I'm definitely using EGL. The time to render a single scene is fine, it's my method for render all 4096 scenes that's slow.

My current sim code:

class Render:
    def __init__(self, model):
        self.model = model
        self.renderer = mujoco.Renderer(model)

    def __call__(self, rng, batch):
        data = mjx.get_data(self.model, batch)
        for i, d in enumerate(data):
            self.renderer.update_scene(d)
            pixels = self.renderer.render()
            dataset.image(i, d.time, pixels)
        return rng, batch

When I have a lot of parallel sims, the copy and for loop can take a good amount of time - I can mitigate this slightly using a background thread especially if I'm simulating at 1 khz and rendering at 10 hz. Is there a way I can avoid doing this copy and loop? It feels like it should be possible since the data is on the GPU which is doing the rendering.

Model:

<mujoco>
  <option timestep=".001">
  </option>
  <asset>
    <texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3" rgb2=".2 .3 .4" width="640" height="480"/>
    <material name="grid" texture="grid" texrepeat="8 8" reflectance=".2"/>
  </asset>
  <worldbody>
    <geom name="floor" type="plane" pos="0 0 -5" size="20 20 .1" material="grid"/>
    <light pos="0 -.4 1"/>
    <body name="oleg" pos="0 0 0">
      <body name="thigh" pos="0 0 0">
        <joint type="slide" axis="1 0 0" pos="0 0 0" limited="true" range="-10 10"/>
        <joint type="slide" axis="0 0 1" pos="0 0 0" limited="true" range="-10 0"/>
        <joint name="thigh_joint" type="hinge" axis="0 -1 0" pos="0 0 6" limited="true" range="-60 60"/>
        <joint type="hinge" axis="-1 0 0" pos="0 0 6" limited="true" range="-5 5"/>
        <geom pos="0 0 4" type="box" size="0.2 0.1 2.0" rgba="1 1 0 1" mass="0.5"/>
        <body name="calf" pos="0 0 0">
          <joint name="knee_joint" axis="0 -1 0" pos="0 0 2"/>
          <geom pos="0 0 0" type="box" size="0.2 0.1 2.0" rgba="1 0 0 1" mass="0.3"/>
        </body>
      </body>
    </body>
  </worldbody>
  <actuator>
    <motor name='thigh_motor' ctrllimited="true" ctrlrange="-100.0 100.0" joint="thigh_joint"/>
    <motor name='knee_motor' ctrllimited="true" ctrlrange="-100.0 100.0" joint="knee_joint"/>
  </actuator>
</mujoco>
@c0g c0g added the question Request for help or information label Apr 18, 2024
@erikfrey
Copy link
Collaborator

This is a good question. You're right - the culprit here is mjx.get_data - any scenario where you are copying data to host and then back to GPU to render is going to kill throughput.

Doing this right requires building interfaces to allow MJX and a renderer to share the same GPU buffers. The most promising avenue for this is Madrona, and we are experimenting with them on ways to integrate to MJX, but it's early days.

We are also working on an simple depth renderer that would be native to MJX using mjx.ray, however I suspect the Madrona integration to be both faster and more powerful. Stay tuned, this is actively under development.

@c0g
Copy link
Author

c0g commented May 3, 2024

Thanks for the answer :-) glad to hear it’s being worked on. Anything early I can play with?

@Andrew-Luo1
Copy link
Contributor

I have a depth camera implementation using mjx.ray from a personal project, that I discuss here. It renders quite quickly and I've successfully used it for vision-based RL.

@btaba btaba mentioned this issue May 21, 2024
Open
@erikfrey erikfrey mentioned this issue May 22, 2024
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
question Request for help or information
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
3 participants
@erikfrey @c0g @Andrew-Luo1