Simulate Images for ML in PyBullet — The Quick & Easy Way | by Mason McGough | Towards Data Science

When applying deep Reinforcement Learning (RL) to robotics, we are faced with a conundrum: how do we train a robot to do a task when deep learning requires hundreds of thousands, even millions, of examples? To achieve 96% grasp success on never-before-seen objects, researchers at Google and Berkeley trained a robotic agent through 580,000 real-world grasp attempts. This feat took seven robots and several weeks to accomplish. Without Google resources, it may seem hopeless for the average ML practitioner. We cannot expect to easily run hundreds of thousands of iterations of training using a physical robot, which is subject to wear-and-tear and requires human supervision, neither of which comes cheap. It would be much more feasible if we could pretrain such RL algorithms to drastically reduce the number of real world attempts needed.
With the advent of deep learning, RL techniques have matured but so has the demand for data. In an attempt to bridge this gap, many researchers are exploring the synthetic generation of training data, utilizing 3D rendering techniques to produce mock-ups of the task environment. While this technique works wonders in the simulated environment, it does not generalize well to the real world. Anyone who has worked closely with deep learning has learned how effective it is at exploiting nuances in training data to “cheat” a task. In practice, the real world is out-of-sample for these models trained in simulation and, unsurprisingly, they fail.
Recent research is focused towards reducing the “reality gap,” those noticeable differences between the real world and 3D-rendered facsimiles, to pretrain robotic agents in simulation and thus drastically reduce the amount of real-world training needed. By aggressively randomizing the appearance and dynamics of a simulation, models learn features that in theory should generalize to the real world. This can even reduce the number of training sessions needed to train a robot by 99% according to recent research. For computer vision tasks where producing hundreds of thousands, or even millions, of training images is infeasible, this method is highly appealing.
Let’s say for example we wish to generate synthetic training data to pretrain a classifier that can distinguish between a collection of small objects. We wish to select an object, drop it into a virtual “bin,” and render a screenshot. Ideally, we want to repeat this process thousands of times with randomization to create a rich dataset of images of objects. To this end (and inspired by the work done at OpenAI and Google), I recently began my journey to generate training data for my own applications. This led me to PyBullet, a Python module designed for robotics and machine learning applications based on the Bullet Physics SDK.
Unlike other 3D rendering solutions like Maya or Blender, PyBullet is focused on robotics and has native implementations for concepts like joints, dynamics simulation, forward and inverse kinematics, and more. In addition, it can be easily installed with a package manager, allowing you to seamlessly integrate other Python packages, like NumPy and TensorFlow, into your simulations. While similar to proprietary physics simulation software MuJoCo, PyBullet is both free and easy to install, so it is a great choice for someone looking to experiment with simulation and robotics. This guide is meant as an introduction for those wishing to generate training images, but the official PyBullet quick start guide can be found here. Note that this tutorial assumes you are using Python version ≥3.6 (if you are using an older version, try creating a virtual environment for this tutorial).
First things first, we will install PyBullet to our Python environment.
pip install numpypip install pybullet
Note that installing NumPy first is optional, but recommended when rendering images, due to the overhead of copying image buffers between C/C++ and Python.
Now that you have everything successfully installed, let us dive right into a PyBullet session. PyBullet relies on a client-server model, where your Python session sends and receives data from a simulation server. There are chiefly two types of servers: DIRECT and GUI. As the name implies, GUI allows you to see the physical simulation in a GUI. DIRECT is useful for rendering headlessly and can be used for efficient rendering without a GPU. DIRECT will be the mode of choice when we are ready to render thousands of images as training data, but for now we will select GUI. Let us begin by initializing a PyBullet server. Start by launching a python interactive session and then enter the following:
import pybullet as pbphysicsClient = pb.connect(pb.GUI)
This step is necessary, as it instantiates the world simulation that you will work with. The physicsClient variable holds a unique ID for that server. This makes it possible in PyBullet to run multiple servers, even across multiple machines, and manage them from a single script by keeping track of these IDs. Many entities in PyBullet are similarly referred to by their IDs, including objects, colliders, and textures.
Now you should see a simulation window pop up. That’s great! It means your simulation is working. You should see a primary viewport with additional windows for RGB, depth, and segmentation data. The world should be completely empty. Let’s first create a plane for our simulation. This plane comes with the pybullet_data package and can be easily loaded into any simulation.
import pybullet_datapb.setAdditionalSearchPath(pybullet_data.getDataPath())planeId = pb.loadURDF(‘plane.urdf’)
This does a few things. First, it loads another module called pybullet_data, which contains placeholder URDFs for many useful things. This is added to the path, so that when we finally call loadURDF it knows where to find ‘plane.urdf’ and loads it into our scene. Now we should have a plane that spans our scene.
The most supported file formats in PyBullet are URDF, SDF, and MJCF. These formats are very easy to load and configure and have their own specialized loading functions. However, many of the 3D models we find on the web are in the simple and popular Wavefront OBJ file format. It is possible to load these, but requires knowing a bit more about the nuts-and-bolts of models in PyBullet, so we will bite the bullet and learn to load our own models with OBJ files.
We will be using a dataset of procedurally generated objects for simulation tasks. This set consists of 1000 OBJ files which include corresponding colliders, materials, and URDF versions. Download those here and extract them into a folder called “random_urdfs” in your project folder.
Every typical 3D object in a physical simulation consists of a visual mesh to describe the way an object appears and a collider mesh that describes an object’s physical dimensions and how it will interact with the world, such as how it collides with other physical objects. Collider meshes, or simply colliders, are typically simplified representations of our visual mesh that are computationally much cheaper to calculate collisions with. Using the two types of meshes together enables our physical simulations to run efficiently while still being visually appealing.
Let us start by loading the visual shape for an object using the createVisualShape function. We can start with an object from our procedurally generated objects dataset.
visualShapeId = pb.createVisualShape(    shapeType=pb.GEOM_MESH,    fileName='random_urdfs/000/000.obj',    rgbaColor=None,    meshScale=[0.1, 0.1, 0.1])
The first parameter shapeType tells PyBullet what kind of shape we are loading. These include spheres, boxes, cylinders, planes, capsules, and meshes. In our case, we wish to load a custom mesh described in our OBJ file. You can also describe the RGB color (as a length-4 list describing red, green, blue, and alpha channels), which we set to None here. We will replace the object’s texture later. Lastly, OBJ files do not intrinsically have units, so we may need to specify a mesh scale to ensure that our object is not too large or too small.
Now we have a visual shape to reference for our object, but no physical shape. Let us load this using the createCollisionShape function.
collisionShapeId = pb.createCollisionShape(    shapeType=pb.GEOM_MESH,    fileName='random_urdfs/000/000_coll.obj',    meshScale=[0.1, 0.1, 0.1])
Similar to the createVisualShape function, we need to specify that we want a mesh using shapeType and where to find the mesh data with fileName. We can and should also specify the mesh scale so that the collider has the same shape as the visual shape.
At this point, we have references to a visual shape and a collider, but the two aren’t connected and we have no object yet. That is where createMultiBody comes in. This function unites these two and creates an instance of an object.
multiBodyId = pb.createMultiBody(    baseMass=1.0,    baseCollisionShapeIndex=collisionShapeId,     baseVisualShapeIndex=visualShapeId,    basePosition=[0, 0, 1],    baseOrientation=pb.getQuaternionFromEuler([0, 0, 0]))
We pass the IDs of our collider and visual shapes to createMultiBody to create a new object. This object needs a position and orientation to know where to place it in the physical world and a mass to know how it should interact with other objects. Set position with a length-3 list of x, y, z coordinates. For simplicity, we set mass to 1.0 for now.
Orientation is a little trickier and requires that we specify it in a length-4 list of quaternions. Quaternions are a useful 4-dimensional representation for rotation in 3D spaces that are crucially used in computer graphics software like Blender, Maya, and Unity to avoid rotational artifacts such as gimbal lock. PyBullet has a convenience function called getQuaternionFromEuler that allows us to describe rotation in rotations about the x, y, and z axes, which are easier for human beings to visualize and understand. This representation of rotation is called Euler angles.
Note: if you receive a message like “Unsupported texture image format [random_urdfs/000/000.png],” then ignore it. We will override the default texture in the next step anyways.
Success! We have generated a physical object and placed it in our world. Recall from earlier that we want to randomize the appearance and dynamics of our 3D models. Let us start with the appearance and download and extract the Describable Textures Dataset here. This contains thousands of images comprised of 47 texture categories. Extract these into a folder called “dtd”. Let us pick one of these textures at random and apply it to our model.
import os, glob, randomtexture_paths = glob.glob(os.path.join('dtd', '**', '*.jpg'), recursive=True)random_texture_path = texture_paths[random.randint(0, len(texture_paths) - 1)]textureId = pb.loadTexture(random_texture_path)pb.changeVisualShape(multiBodyId, -1, textureUniqueId=textureId)
This snippet recursively grabs a list of paths to all of the textures in folder “dtd” with extension “.jpg” and randomly samples a single one. Then it passes it to function loadTexture to load a texture and create an ID. It then uses the changeVisualShape function to apply the texture with textureId to object with multiBodyId. Now your object should have the texture applied to it. You can also apply this same method to attach a texture to the plane if you wish.
Note that the -1 is necessary for required parameter jointIndex, which specifies the link you wish to apply the texture to. Links are useful for creating objects with hierarchies of moving parts, like robots with many joints. This is an important concept in PyBullet, but is out of the scope of this tutorial, so we supply a value “-1” to tell it to apply it to the base object rather than any link.
So far, we have created our simulation, but nothing is moving. There are a few reasons for this. First of all, we haven’t applied any gravity to the scene, so our object is floating in space. Let us do that real quick.
pb.setGravity(0, 0, -9.8)
That applies a downward negative acceleration of -9.8 m/s2 to our objects in the scene, just like we experience on the surface of planet Earth. However, your object still is not falling like you would expect. That is because the simulation is not stepping.
There are two ways to step through your simulation: manually call stepSimulation to proceed the simulation by one time step or call setRealTimeSimulation to render in real time. For demonstration purposes, we will use the latter, but you will probably wish to use stepSimulation once we start rendering batches of objects in a loop.
pb.setRealTimeSimulation(1)
Set to “1” to enable or “0” to disable. Now your object should fall and hit the ground. By the way, the default time step is 1/240 seconds, but you can override this default if desired using the setTimeStep function.
We now have a method to load a plane, place an OBJ file on our plane, apply a texture to the object, and render the scene. The last piece of the puzzle is rendering an image of our object. Essentially, we want to place a virtual camera over our object, point it downward at our object, and take a picture.
Whether we are looking to solve a classification or segmentation task, PyBullet makes taking pictures of our simulation fairly straightforward with the getCameraImage function, which renders an RGB image, a segmentation mask, and a depth buffer simultaneously. However, we need to first specify the extrinsic and intrinsic properties of our camera, a process that is a bit more involved than it may seem.
The camera extrinsic properties essentially describe how our camera is positioned in the world, such as its orientation and position. The camera intrinsic properties describe properties of the camera itself, such as the field-of-view (FOV) and the aspect ratio of its sensor (e.g. 4:3, 16:9, etc.). In PyBullet, we describe these properties with a view matrix and a projection matrix, respectively. The PyBullet quick start guide links to a great resource to learn more about the camera extrinsic and intrinsic properties and their corresponding matrices.
The camera view matrix is a complicated 4x4 matrix, but in simplest terms it describes where the camera is and in what direction it is pointing. There are some useful helper functions for creating this matrix by specifying position and rotation more directly. The function computeViewMatrix can create this matrix in exchange for three vectors.
viewMatrix = pb.computeViewMatrix(    cameraEyePosition=[0, 0, 3],    cameraTargetPosition=[0, 0, 0],    cameraUpVector=[0, 1, 0])
As the name implies, cameraEyePosition describes the physical location of the camera in x, y, and z coordinates. We will place it 3 units directly above the plane, centered over our object. Next we describe the point that we wish the camera to face with cameraTargetPosition. We want it to look straight down at our object, so we point it at the origin. Lastly, we need to describe the orientation of the camera. We do this with cameraUpVector, which is a vector that points out of the top of our camera. We point this out along the y axis. Our camera should now be pointing directly downward so that the image plane that it renders is perfectly parallel with the plane we generated.
The projection matrix, much like the view matrix, can be created using a few helper functions. In this case, the computeProjectionMatrixFOV function describes our camera’s intrinsic properties in the simplest and most pertinent ways to our use case.
projectionMatrix = pb.computeProjectionMatrixFOV(    fov=45.0,    aspect=1.0,    nearVal=0.1,    farVal=3.1)
The fov parameter specifies the camera FOV in degrees. It describes how “wide” the camera’s visual field is and is similar in principle to the focal length of a real camera. The aspect parameter describes the camera aspect ratio, which we set to 1.0 here for simplicity. The nearVal and farVal parameters describe the minimum and maximum distance, respectively, at which the camera will render objects. Since our camera is 3 units above the plane, we set farVal to 3.1 to just include that plane. These parameters are defined in the OpenGL documentation and you can read more about them here.
We have described our camera’s properties and are now ready to render some images. Let us do that with getCameraImage.
This function returns three image buffers: rgbImg, depthImg, and segImg. rgbImg is a uint8 image with red, green, blue, and alpha channels of the camera’s visuals. depthImg is a floating-point grayscale image that describes the distance of individual rendered pixels from the camera. It can be used to model the field-of-view of a real-world depth sensor, such as a Microsoft Kinect. Lastly, segImg is a segmentation mask of the image where pixels each contain unique integers with object IDs. These are invaluable for training segmentation algorithms for robotic agents, such as a robotic arm that needs to identify objects to sort into respective bins or for a driverless car that wants to identify pedestrians, street signs, and roads.
width, height, rgbImg, depthImg, segImg = pb.getCameraImage(    width=224,     height=224,    viewMatrix=viewMatrix,    projectionMatrix=projectionMatrix)
Once you call this function, you should see the images that you rendered in the three windows in your viewer. If that is the case, then congrats! You have rendered your first object. You would probably want to convert these image buffers to NumPy vectors and save them with an imaging library such as PIL or imageio.
PyBullet makes our task of creating a virtual scene to generate training images fairly simple. With the concepts discussed in this tutorial, you should be able to generate your own datasets for classification or segmentation. Additional augmentations might be worth exploring, such as adding noise to the position and orientation of the camera and rendered objects. One might also wish to replace the plane with a mock-up of the real-world environment you wish to train for. Hopefully this tutorial is a solid launching point for you to develop your own physical simulations for machine learning.