Visualization of Physics Simulation in Google Colab

This article is for personal use to archive the visulization method for physics simulation in Google Colab.

Caution

Before we begin, I strongly recommend you just use your personal computer environment to visualize the physics engine. There is no "practical" way to render the physics engine in Colab, and modern computers are very good at rendering thanks to the game community. The reason I am writing is that there are situations where you do not want to mess up the computer or want to train a reinforcement learning model with vision, which requires a lot of computation resources.

Why visualization?

The simplest and most intuitive way to know whether the simulation outcome has come out as expected is to "see" the result. For example, if you want to know if your pick-and-place algorithm for robot manipulators is good or not, you can visualize the physics engine and see if robots move objects to the target. Therefore, all physics simulation engines provide basic rendering options to users, using OpenGL or gaming engines like Unity or Unreal. However, if you use an IPython environment like Jupyter Notebook or Colab, these rendering options usually do not work. Therefore, in this article, we look for possible ways to visualize physics engine results. Of course, since we are using an IPython environment, we only consider physics engines like Pybullet or Mujuco, not C++ based engines like Gazebo.

1. Remote desktop (policy violation)

Before Google disallowed it for Colab in 2021, remote desktop was the best method to visualize the physics engine in Colab. There are three steps to doing this: setup SSH server in the Colab instance using service like ngrok, install VirtualGL in the Colab instance so that it can run OpenGL programs, and access the Colab instance using a VNC program. However, this is not legal anymore and Google will stop you if you try this.

2. Recording video

This is the most common way that physics engine tutorials introduce this. There are three steps to doing this: saving all frames of the physics engine for a predetermined time duration; converting the frames to a video; and playing the video in the IPython environment. And there are two ways to implement this: using HTML and Matplotlib, or using the Numpng library to make animated PNG (APNG). The HTML method is complicated to my personal standard as I do not have much experience with websites. Therefore, I attached a code for the APNG method. However, if you are interested in the HTML method, go to this Mujoco tutorial link.

from IPython.display import Image
from numpngw import write_apng

frames = []
duration = 100 # custom time duration

#physics_environment is a pseudo physics engine
for physics_environment.time < duration:
    physics_environment.step()
    frame = physics_environment.render()
    frames.append(frame)

delay = 1000/physics_environment.fps # fps: frames per second
write_apng('result.png', frames, delay=delay) # delay in ms
Image(filename='result.png')

The downside of this method is that you have to wait until the physics simulation ends. This is critical because we usually do not know how long it will take for a robot to carry out a task at first, which is the duration in the code. Also, there might be cases where we want to terminate the physics engine earlier because the visual results do not seem good.

3. Live video

I prefer live visualization since I frequently terminate the physics engine while it is being executed to debug errors. The way to do this is to keep printing the new images and deleting old images in the IPython environment.

from google.colab.patches import cv2_imshow
from IPython.display import clear_output

frames = []
duration = 100 # custom time duration

#physics_environment is a pseudo physics engine
for physics_environment.time < duration:
    physics_environment.step()
    frame = physics_environment.render()
    cv2_imshow(frame)
    clear_output(wait=True)

However, this creates a blinking effect because of the time interval between the deletion of the old image and the printing of the image. Therefore, rather than deleting the previous image, we update the previous image to the new image, using PIL, as below.

from IPython.display import display, Image
from PIL import Image

display_handle=display(None, display_id=True)

frames = []
duration = 100 # custom time duration

#physics_environment is a pseudo physics engine
for physics_environment.time < duration:
    physics_environment.step()
    frame = physics_environment.render()
    # 'RGB' for Mujoco and 'RGBA" for Pybullet
    display_handle.update(Image.fromarray(frame, 'RGB'))

Then, we can get uninterrupted video of the physics engine. The downside of this is that, since the physics engine time does not match with real-world time, the time of the simulation result is kind of unpredictable. To resolve this, you can either code to output a timer along with the video, or you can code to force the set FPS of the video. However, the fact that this method is slow cannot be solved. The below code is an example code using Pybullet, which computes a frame every 10 steps to match the real-world time.

# Install the PyBullet module.
!pip install -U pybullet

import pybullet as p
import time
import pybullet_data

from IPython.display import display
from PIL import Image

physicsClient = p.connect(p.DIRECT)
p.setAdditionalSearchPath(pybullet_data.getDataPath())

p.setGravity(0,0,-10)
planeId = p.loadURDF('plane.urdf')
robotStartPos = [0, 0, 1]
robotStartOrientation = p.getQuaternionFromEuler([0,0,0])
# robots will explode because they are in the same position and collide extremely
robotId1 = p.loadURDF('r2d2.urdf', robotStartPos, robotStartOrientation)
robotId2 = p.loadURDF('r2d2.urdf', robotStartPos, robotStartOrientation)
robotId3 = p.loadURDF('r2d2.urdf', robotStartPos, robotStartOrientation)

display_handle=display(None, display_id=True)

for i in range(10000):
  p.stepSimulation()

  width = 640
  height = 480
  # computing the image once every 10 steps
  if i%10 == 9:
    img = p.getCameraImage(
        width,
        height,
        viewMatrix=p.computeViewMatrixFromYawPitchRoll(
            cameraTargetPosition=[0, 0, 0],
            distance=4,
            yaw=60,
            pitch=-10,
            roll=0,
            upAxisIndex=2,
        ),
        projectionMatrix=p.computeProjectionMatrixFOV(
            fov=60,
            aspect=width/height,
            nearVal=0.01,
            farVal=100,
        ),
        shadow=False,
        lightDirection=[1, 1, 1],
        renderer=p.ER_BULLET_HARDWARE_OPENGL
    )

    width, height, rgba, depth, mask = img
    display_handle.update(Image.fromarray(rgba, 'RGBA'))

p.disconnect()