BulletArm: An Open-Source Robotic Manipulation Benchmark and Learning Framework

The design philosophy behind our framework focuses on four key principles:

1. Reproducibility: A key challenge when developing new learning algorithms is the difficulty in comparing them to previous work. In robotics, this problem is especially prevalent as different researchers have drastically different robotic setups. This can range from small differences, such as workspace size or degradation of objects, to large differences such as the robot used to preform the experiments. Moving to simulation allows for the standardization of these factors but can impact the performance of the trained algorithm in the real-world. We aim to encourage more direct comparisons between works by providing a flexible simulation environment and a number of baselines to compare against.

2. Extensibility: Although we include a number of tasks, control types, and robots; there will always be a need for additional development in these areas. Using our framework, users can easily add new tasks, robots, and objects. We make the choice to not restrict tasks, allowing users more freedom create interesting domains. Figure 1 shows an example of creating a new task using our framework.

3. Performance: Deep learning methods are often time consuming, slow processes and the addition of a physics simulator can lead to long training times. We have spent a significant portion of time in ensuring that our framework will not bottleneck training by optimizing the simulations and allowing the user to run many environments in parallel.

4. Usability: A good open-source framework should be easy to use and understand. We provide extensive documentation detailing the key components of our framework and a set of tutorials demonstrating both how to use the environments and how to extend them.

Our simulation setup (Figure 5) consists of a robot arm mounted on the floor of the simulation environment, a workspace in front of the robot where objects are generated, and a sensor. Typically, we use top-down sensors which generate heightmaps of the workspace. As we restrict the perception data to only the defined workspace, we choose to not add unnecessary elements to the setup such as a table for the arm to sit upon. Currently there are four different robot arms available in BulletArm (Figure 3): KUKA IIWA, Frane Emika Panda, Universal Robots UR5 with either a simple parallel jaw gripper or the Robotiq 2F-85 gripper.

Environment, Configuration, and Episode are three key terms within our framework. An environment is an instance of the PyBullet simulator in which the robot interacts with objects while trying to solve some task. This includes the initial environment state, the reward function, and termination requirements. A configuration contains additional specifications for the task such as the robotic arm, the size of the workspace, the physics mode, etc (see Appendix B for an full list of parameters). Episodes are generated by taking actions (steps) within an environment until the episode ends. An episode trajectory $\tau$ contains a series of observations $o$, actions $a$, and rewards $r$: $\tau=[(o_{0},a_{0},r_{0}),...,(o_{T},a_{T},r_{T})]$.

Users interface with the learning environment through the EnvironmentFactory and EnvironmentRunner classes. The EnvironmentFactory is the entry point and creates the Environment class specified by the Configuration passed as input. The EnvironmentFactory can create either a single environment or multiple environmaents meant to be run in parallel. In either case, an EnvironmentRunner instance is returned and provides the API which interacts with the environments. This API, Figure 1, is modelled after the typical agent-environment RL setup popularized by OpenAI Gym [4].

The benchmark tasks we provide have a sparse reward function which returns $+1$ on successful task completion and $0$ otherwise. While we find this reward function to be advantageous as it avoids problems due to reward shaping, we do not require that new tasks conform to this. When defining a new task, the reward function defaults to sparse but users can easily define their custom reward for a new task. We separate our tasks into two categories based on the action spaces: open-loop control and closed-loop control. These two control modes are commonly used in robotics manipulation research.

Expert demonstrations are crucial to many robotic learning fields. Methods such as imitation and model-based learning, for example, learn directly from expert demonstrations. Additionally, we find that in the context of reinforcement learning, it is vital to seed learning with expert demonstrations due to the difficulties in exploring large state-action spaces. We provide two types of planners to facilitate expert data generation: the Waypoint Planner and the Deconstruction Planner.

The Waypoint Planner is a online planning method which moves the end-effector through a series of waypoints in the workspace. We define a waypoint as $w_{t}=(p_{t},a_{t})$ where $p_{t}$ is the desired pose of the end effector and $a_{t}$ is the action primitive to execute at that pose. These waypoints can either be absolute positions in the workspace or positions relative to the objects in the workspace. In open-loop control, the planner returns the waypoint $w_{t}$ as the action executed at time $t$. In close-loop control, the planner will continuously return a small displacement from the current end-effector pose as the action at time $t$. This process is repeated until the waypoint has been reached. The Deconstruction Planner is a more limited planning method which can only be applied to pick-and-place tasks where the goal is to arrange objects in a specific manner. For example, we utilize this planner for the various block construction tasks examined in this work (Figure 2). When using this planner, the workspace is initialized with the objects in their target configuration and objects are then removed one-by-one until the initial state of the task is reached. This deconstruction trajectory, is then reversed to produce an expert construction trajectory, $\tau_{expert}=\textit{reverse}(\tau_{deconstruct})$.

In the open-loop environments, the agent controls the target pose of the end-effector, resulting in a shorter time horizon for complex tasks. The open-loop environment is comprised of a robot arm, a workspace, and a camera above the workspace providing the observation (Figure 5). The action space is defined as the cross product of the gripper motion $A_{g}=\{\textsc{pick},\textsc{place}\}$ and the target pose of the gripper for that motion $A_{p}$, $A=A_{g}\times A_{p}$. The state space is defined as $s=(I,H,g)\in S$ (Figure 5), where $I$ is a top-down heightmap of the workspace; $g\in\{\textsc{holding},\textsc{empty}\}$ denotes the gripper state; and $H$ is an in-hand image that shows the object currently being held. If the last action is pick, then $H$ is a crop of the heightmap in the previous time step centered at the pick position. If the last action is place, $H$ is set to a zero value image.

BulletArm provides three different action spaces for $A_{p}$: $A_{p}\in\{A^{xy},A^{xy\theta},A^{\mathrm{SE}(3)}\}$. The first option $(x,y)\in A^{xy}$ only controls the $(x,y)$ components of the gripper pose, where the rotation of the gripper is fixed; and $z$ is selected using a heuristic function that first reads the maximal height in the region around $(x,y)$ and then either adds a positive offset for a place action or a negative offset for pick action. The second option $(x,y,\theta)\in A^{xy\theta}$ adds control of the rotation $\theta$ along the $z$-axis. The third option $(x,y,z,\theta,\phi,\psi)\in A^{\mathrm{SE}(3)}$ controls the full 6 degree of freedom pose of the gripper, including the rotation along the $y$-axis $\phi$ and the rotation along the $x$-axis $\psi$. $A^{xy}$ and $A^{xy\theta}$ are suited for tasks that only require top-down manipulations, while $A^{\mathrm{SE}(3)}$ is designed for solving complex tasks involving out-of-plane rotations. The definition of the state space and the action space in the open-loop environments also enables effortless sim2real transfer. One can reproduce the observation in Figure 5 in the real-world using an overhead depth camera and transfer the learned policy [43, 45].

Figure 6 shows the 12 basic open-loop environments that can be solved using top-down actions. Those environments can be categorized into two collections, a set of block structure tasks (Figures 6a-6g), and a set of more realistic tasks (Figures 6h and 6l). The block structure tasks require the robot to build different goal structures using blocks. The more realistic tasks require the robot to finish some real-world problems, for example, arranging bottles or supervising Covid tests. We use the default sparse reward function for all open-loop environments, i.e., +1 reward for reaching the goal, and 0 otherwise. See Appendix A.1 for a detailed description of the tasks.

The close-loop environments require the agent to control the delta pose of the end-effector, allowing the agent more control and enabling us to solve more contact-rich tasks. These environments have a similar setup to the open-loop domain but to avoid the occlusion caused by the arm, we instead use two side-view cameras pointing to the workspace (Figure 8a). The heightmap $I$ is generated by first acquiring a fused point cloud from the two cameras (Figure 8b) and then performing an orthographic projection (Figure 8c). This orthographic projection is centered with respect to the gripper. In practice, this process can be replaced by putting a simulated orthographic camera at the position of the gripper to speed up the simulation. The state space is defined as a tuple $s=(I,g)\in S$, where $g\in\{\textsc{holding},\textsc{empty}\}$ is the gripper state indicating if there is an object being held by gripper. The action space is defined as the cross product of the gripper open width $A_{\lambda}$ and the delta motion of the gripper $A_{\delta}$, $A=A_{\lambda}\times A_{\delta}$.

Two different action spaces are available for $A_{\delta}$: $A_{\delta}\in\{A_{\delta}^{xyz},A_{\delta}^{xyz\theta}\}$. In $A_{\delta}^{xyz}$, the robot controls the change of the $x,y,z$ position of the gripper, where the top-down orientation $\theta$ is fixed. In $A_{\delta}^{xyz\theta}$, the robot controls the change of the $x,y,z$ position and the top-down orientation $\theta$ of the gripper. Figure 9 shows the 10 close-loop environments. We provide a default sparse reward function for all environments. See Appendix A.2 for a detailed description of the tasks.

In the Open-Loop 3D Benchmark, the agent needs to solve the open-loop tasks shown in Figure 6 using the $A_{g}\times A^{xy\theta}$ action space (see Section 4.1). We provide a set of baseline algorithms that explicitly control $(x,y,\theta)\in A^{xy\theta}$ and select the gripper motion using the following heuristic: a pick action will be executed if $g=\textsc{empty}$ and a place action will be executed if $g=\textsc{holding}$. The baselines include: (1) DQN [30], (2) ADET [24], (3) DQfD [15], and (4) SDQfD [43]. The network architectures for these different methods can be used interchangeably. We provide the following network architectures:

1. 1.

   CNN ASR [43]: A two-hierarchy architecture that selects $(x,y)$ and $\theta$ sequentially.

2. 2.

   Equivariant ASR (Equi ASR) [45]: Similar to ASR, but instead of using conventional CNNs, equivariant steerable CNNs [5, 46] are used to capture the rotation symmetry of the tasks.

3. 3.

   FCN: a Fully Convolutional Network (FCN) [29] which outputs a $n$ channel action-value map for each discrete rotation.

4. 4.

   Equivariant FCN [45]: Similar to FCN, but instead of using conventional CNNs, equivariant steerable CNNs are used.

5. 5.

   Rot FCN [52, 51]: A FCN with 1-channel input and output, the rotation is encoded by rotating the input and output for each $\theta$.

In this section, we show the performance of SDQfD (which is shown to be better than DQN, ADET, and DQfD [43]. See the performance of DQN, ADET and DQfD in Appendix E) equipped with CNN ASR, Equi ASR, FCN, and Rot FCN. We evaluate SDQfD in the 12 environments shown in Figure 6. Table 2 shows the number of objects, the optimal number of steps per episode, and the max number of steps per episode in the open-loop benchmark experiments. Before the start of training, 200 (500 for Covid Test) episodes of expert data are populated in the replay buffer. Figure 10 shows the results. Equivariant ASR (blue) has the best performance across all environments, then Rot FCN (green) and CNN ASR (red), and finally FCN (purple). Notice that Equivariant ASR is the only method that is capable of solving the most challenging tasks (e.g., Improvise House Building 3 and Covid Test).

The Close-Loop 4D Benchmark requires the agent to solve the close-loop tasks shown in Figure 9 in the 5-dimensional action space of $(\lambda,x,y,z,\theta)\in A_{\lambda}\times A_{\delta}^{xyz\theta}\subset\mathbb{R}^{5}$, where the agent controls the positional displacement of the gripper ($x,y,z$), the rotational displacement of the gripper along the $z$ axis ($\theta$), and the open width of the gripper ($\lambda$). We provide the following baseline algorithms: (1) SAC [13], (2) Equivariant SAC [44], (3) RAD [25] SAC: SAC with data augmentation, (4) DrQ [23] SAC: Similar to (3), but performs data augmentation when calculating the $Q$-target and the loss, and (5) FERM [53]: A Combination of SAC and contrastive learning [26] using data augmentation. Additionally, we also provide a variation of SAC, SACfD [44], that applies an auxiliary L2 loss towards the expert action to the actor network. SACfD can also be used in combination with (2), (3), and (4) to form Equivariant SACfD, RAD SACfD, DrQ SACfD, and FERM SACfD.

In this section, we show the performance of SACfD, Equivariant SACfD (Equi SACfD), Equivariant SACfD using Prioritized Experience Replay (PER [35]) and data augmentation (Equi SACfD + PER + Aug), and FERM SACfD. (See Appendix F for the performance of RAD SACfD and DrQ SACfD.) We use a continuous action space where $x,y,z\in[-0.05m,0.05m],\theta\in[-\frac{\pi}{4},\frac{\pi}{4}],\lambda\in[0,1]$. We evaluate the various methods in the 10 environments shown in Figure 9. Table 3 shows the number of objects, the optimal number of steps for solving each task, and the maximal number of steps for each episode. In all tasks, we pre-load 100 episodes of expert demonstrations in the replay buffer.

Figure 11 shows the performance of the baselines. Equivariant SACfD with PER and data augmentation (blue) has the best overall performance followed by standard Equivariant SACfD (red). The equivariant algorithms show a significant improvement when compared to the other algorithms which do not encode rotation symmetry, i.e. CNN SACfD and FERM SACfD.