3D reconstruction based on hierarchical reinforcement learning with transferability

3.1Framework overview

As shown in Fig. 1, the goal is to reconstruct a 3D shape as similar to the target 3D model as possible from the provided target image. First, obtain the coarse-grained mesh through heuristic [36] or deep learning methods [1], and then add loops on it. Second, the agent is trained to choose an editing action at each step to edit the mesh. This paper focuses on better training of the agent. The intersection over union (IoU) between the edited mesh and the target 3D model is calculated at each step. If the IoU increases after editing, the agent’s decision is encouraged; otherwise, it will be punished. The agent is constantly trained through trial and error.

As shown in Fig. 2a, the agent observes the current state st at time step t and selects an action at. After the environment executes the action, it obtains the next state st+1 and the reward rt+1. Then, we set st=st+1, rt=rt+1 and repeat the above process until reaching the termination state. The reward will guide the agent to make the decision with the highest discounted cumulative reward R, where R=∑i=0∞γi⁢ri and γ is the discounted factor. The discount factor γ determines the impact of future rewards on the current value: the value of rewards received in the future k time steps is worth only γk times the value of rewards received immediately. The detailed definitions of the state, action and reward can be seen in Section 3.2.

Figure 2b describes the option hierarchical framework. The agent can be decomposed into the top-agent and a set of sub-agents. The top-agent chooses an option, and then the option-related sub-agent will decide to execute which action. To simplify the training of sub-agents, we proposed ASS-Sub-Agent to replace the set of sub-agents. The details of the option hierarchical framework and ASS-Sub-Agent are described separately in Sections 3.3 and 3.4. Finally, Section 3.5 will show how to transfer knowledge to other classes.

Figure 3.

We assign n loops to the cuboids of coarse-grained mesh. We add n=10 loops to 2 cuboids.

3.2Reinforcement learning modeling

In this subsection, the details of state, action and reward are described.

The proposed approach focuses on adjusting a coarse-grained mesh to have more details according to the provided target image. The coarse-grained mesh is composed of several cuboids, which do not have any loops at first. For further editing, the coarse-grained mesh is subdivided by adding loops [72]. The loop is a rectangle on the cuboid surface, in which the plane is vertical to the longest cuboid side. Here, we assign n loops L in the mesh as Fig. 3, where L={Li}i=0n. A loop Li can be represented by its two corner vertices VL=(xl,yl,zl) and VL′=(xl′,yl′,zl′). The target image, the loops L, and the step are used as state. The actions are defined by the movement of the two corner vertices VL and VL′ as shown in Fig. 4. The movement range is controlled with six parameters -3, -2, -1, 1, 2, 3 on three axis directions. Therefore, the number of total actions is 36⁢n.

Finally, the IoU can be used as a metric of the similarity between the edited mesh and the target 3D model. Accordingly, the reward is the increment of IoU after an action is executed. The use of the metric that is highly relevant to human perception can intrinsically optimize the training of the model [73].

Figure 4.

For one loop, the movements of two corner vertices VL and VL′ are defined as actions.

3.3Hierarchical options framework

To define our problem as one that can be handled hierarchically, we are inspired by the options framework [57, 58], which first chooses an option, and then selects an action based on the current option. For a coarse-grained mesh, the HRLT3D first observes which part is the most different from the target shape, and then decides how to adjust it. In this stage, each loop can represent one part of the original mesh. Therefore, we group the actions operating on one loop into a macro action, equivalently referred to as an option. Therefore, 36n actions can be grouped into n options as the output of the top-agent and 36 actions for each sub-agent. The reason why this can achieve better results is that it decomposes the task and simplifies the action space at different levels. This is consistent with the idea of task decomposition in industry [74]. The top level only needs to focus on which loop needs to be adjusted most, while the low level only needs to focus on the specific operation of the loop. It is the cooperation between global exploration and local exploitation, that can significantly improve accuracy and reliability [75, 76].

The agents are trained by updating the value function Q. When the option ω is started in state st, and then we jump k steps to the state st+k where the option ω terminates. Specifically, the loss function of top-agent can be defined as:

where θΩ is the parameter of top-agent network. Rt+k denotes the cumulative discounted reward over k steps and can be calculated as:

To simplify training, we set k to be a fixed hyper parameter, and the termination function βω always will be 1 if k=1. In particular, it is easy to combine heuristic algorithms or expert experiences to accelerate the learning process of top-agent.

The sub-agents are trained using double deep Q-network (DDQN) [77]. The DDQN using double agents to learn value function, current agent and target agent. For option ω, the loss function of sub-agent is as follows:

where θ and θ′ are parameters of current agent and target agent. The at+1𝑚𝑎𝑥 and U are defined as:

Figure 5.

The triple can be the input of top-agent, which will output an option related with loop index. The features of the loop as the augmentation of state are inputted into ASS-Sub-Agent, which decides how to operate on this loop.

3.4ASS-Sub-Agent

Although the options framework can simplify the action space for each agent, the training parameters are increased since we need to learn the sub-agent for each option. Note that the sub-agents for different options are dealing with the same kind of problem except different loops are to be operated. It is feasible to use an agent to implement the functions of a set of sub-agents. In order to identify which option is being operated on, we propose to train a sub-agent based on an augmented state space (ASS-Sub-Agent) by incorporating the distinctive features of options as part of the state. In this stage, Li is added as an augmented part of the state, which can be the input of ASS-Sub-Agent and i is the loop index selected by the top-agent. The whole network is shown in Fig. 5.

Bacon et al. [58] posed that an augmented state may lead to a larger state space, resulting in a poor learning effect. In contrast, an augmented state can promote our learning and transfer for the following two reasons. First, the training process can be shared by ASS-Sub-Agent, which speeds up learning. Second, the diversity of states makes the agent more general to data of other classes.

3.5Transfer learning based on the hierarchical framework

Transfer learning can promote learning efficiency by using existing knowledge [78, 23]. Most transfer learning based on the hierarchical framework trains a set of sub-agents on source data and then reuse them on target data. The top-agent is retrained for each different data point when to execute which sub-agent. We follow a similar architecture, and the details are shown in the pseudocode of Algorithm 3.5.

Transfer learning based on hierarchical framework[1] Initialize the top-agent πΩ with parameters θΩ randomly. Initialize ASS-Sub-Agent πω with previous trained parameters θωSet t=0Choose ωt based on πΩi=t+1, t+2, …, t+kCollect experience (st,at,rt,st+1) based on πωUpdate θω to minimise the loss LωSet t=t+kCalculate cumulative discounted reward over k steps Rt+kCollect experience (st,ωt,Rt+k,st+k)Update θΩ to minimise the loss LΩtermination

To transfer the trained agent from the source data to the target data, the key step is to initialize the parameters of the top-agent and ASS-Sub-Agent. Consistent with the classic method, the proposed transfer procedure retrains the top-agent through random initialization parameters, and reuses the parameters of ASS-Sub-Agent. The experience (st,at,rt,st+1) is collected based on these agent at each time step t. which is used to update the parameters of ASS-Sub-Agent. When the time step jumps from t to t+k, calculate the cumulative discount reward Rt+k. At this time, the experience (st,ωt,Rt+k,st+k) is used to update the top-agent. Differently, the ωt does not represent the chosen sub-agent, but represents the chosen loop which can be as the augmented part of sate and be operated by ASS-Sub-Agent. Besides, the ASS-Sub-Agent can enhance the generalization and transferability among data of different classes. Industrial scenarios are more complex, so the transferability of agent is more important [79].

4.1Experimental settings

The proposed method is based on RL, so we compare it with two recent RL-based baselines [1, 51], which are published in 2020 and 2022 in this filed. The work [1] is an important first stepping stone towards 3D reconstruction by RL. All the RL-based methods use the framework of DDQN [77]. In addition, the GAN-based methods [22, 26] that are promising in the field of 3D reconstruction are compared. The proposed method is unsupervised, where the target 3D models are used only for calculating the IoU. For fair comparisons, 3D-GAN [22] uses the target 3D models only as real objects for the training of the discriminator. ASTA3D [26] directly reconstructs 3D models from 2D images through a differentiable renderer without requiring 3D target models.

This paper follows the dataset from [1, 51], which is the first public dataset since 2020 for this topic. It has three classes: airplanes, cars, and guitars. Each of class contains 650 shapes, where 600 are used for training and 50 are used for testing. In addition, the Fusion 360 Gallery dataset [80] is also used to verify the effectiveness of the proposed method. The Fusion 360 Gallery dataset has rich 2D and 3D geometric data coming from CAD models. We extract 1000 sets of data, which are the pairs of 2D target images and the target 3D models. Among them, 900 are used for training and 100 are used for testing.

For the parameter settings, we set the discounted factor γ= 0.9 for the top-agent and the ASS-Sub-Agent, hyperparameter k= 1 and loop number n= 10. The larger the loop number is, the higher the reconstruction resolution. However, the accuracy of reconstruction may decrease because the action space increases with the number of loops. The agents choose actions greedily with a probability ϵ of 0.98 in the training stage which is used to balance exploration and exploitation and 1 in the testing stage. The coarse-grained original meshes are obtained through a deep network [1].

All of the experiments are trained and tested under the same hardware and software. The model is trained on a single NVIDIA RTX 3090 GPU , the main memory is 64 GB, and the CPU clock frequency is 2.5–4.4 GHz. The network model is implemented on Python 3.6 and PyTorch 1.9.1. The training of our model takes 15 minutes per epoch for the class airplane, car and guitar and 25 minutes for the Fusion 360 Gallery datasets. We trained 200 epochs for each class.

4.2Reconstruction experiments and analysis

In this subsection, we compare the reconstruction results, which use the cumulative reward R and IoU as evaluation metrics. As shown in the Table 1, the GAN-based methods achieve relatively good results, although they only need target 3D models for the training of discriminator or do not need them. For RL-based methods, the reconstruction performances using only DDQN are very poor due to the huge action space in this problem. The method of Lin et al. [1] obtains higher-quality results by combining DDQN with imitation learning (IL). The reconstruction accuracy of Seiya et al. [51] has declined because they paid more attention on viewpoint-independent reconstruction. Our method obtains the highest cumulative reward and IoU for all of three classes. Moreover, we can achieve better performances even without IL(only DDQN).

To intuitively compare the improvement, we list the cumulative reward improvement of our method over Lin et al. [1], which is the state-of-art method recently in this field. As shown in Table 2, our method improved by 47.4%, 35.7% and 15.2% on airplanes, cars and guitars respectively.

Table 3 shows the standard deviation of the accumulated reward, which can reflect the stability of the training. We can see that our method achieves a smaller standard deviation among the three classes.

Figure 6 displays the qualitative results, where the mesh model reconstructed by our method is visually closer to the target image. As shown in Table 4 and Fig. 7, the quantitative and qualitative results of the Fusion 360 Gallery dataset are still better than those of the baselines. Although the objects in the Fusion 360 Gallery dataset do not belong to the same class, which makes the learning of the neural network more difficult, our method has achieved relatively good results in terms of quantity and quality.

Figure 7.

Qualitative results of our method, RL-based baseline [1] and GAN-based baseline [22] on Fusion 360 Gallery dataset. The first row shows the target images. The second row and the third row show the mesh results of RL-based baseline and the voxel results of GAN-based baseline separately. The last row is our mesh results.

Figure 8.

Comparison of our method with and without transfer on three classes. The charts show the change of cumulative reward with the training epoch. The orange line represents the change of our method with transfer, and blue one is else.

4.3Transferability experiments and analysis

In addition to the performance improvement, in this subsection, we confirm the transferability of HRLT3D. That is, we train the top-agent and the ASS-Sub-Agent on source class data first. Then, we re-train the top-agent and fine-tune the ASS-Sub-Agent on the target class data. Finally, we test the performance on target class data. Specifically, we conduct three groups of experiments compared with the state-of-art RL-based method [1]. We use the following representative metrics [81, 82] to evaluate transfer performances:

At first, we show the impact of training with transfer in Fig. 8. The orange line represents the test result changing with the training epoch, while the blue line is that without transfer. As we can see, agent with transfer can explore better JP, and it can converge to a fairly or even better AP. To demonstrate that our method improves the transferability, we also compared the effects of our method and baseline with transferred knowledge obtained from source class data. Besides, different greedy probabilities will affect the exploration of the environment, and correspondingly affect the transfer results. Therefore, two kinds of greedy probabilities are explored, which are the fixed probability of 0.98 and the incremental probability increasing linearly from 0.5 to 0.98. Table 5 quantifies the transfer learning performances of our method and the baseline under different probabilities. We carried out three groups of experiments, and the source class data for each are listed in brackets. It can be seen from the Table 5 that our method under fixed or incremental greedy probability can achieve better results. We use Mann-Whitney U test to verify whether the proposed method has been significantly improved compared with baseline. The Mann Whitney U Test results for the three evaluation metrics are U= 0, P= 0.01, which means that the proposed method has been significantly improved on the three evaluation metrics with 99.99% probability. As shown in Fig. 9, we can observe the following three phenomena. First, baseline with fixed greedy probability can also improve the JP, but it has a negative impact in the later stage due to the poor generalization of the agent network. Second, transfer learning with incremental greedy probability is generally beneficial to the AP. Third, whether using fixed greedy probability or incremental greedy probability, the proposed method gets better results than the baseline.

Figure 9.

Comparison of the transferability between our method and baseline on three classes. The charts show the change of cumulative reward with the training epoch. The orange lines represent the change of our method with transfer learning, and blue lines represent the change of baseline with transfer learning. The solid lines represents the use of fixed greedy probability, and the dotted lines represent the use of incremental greedy probability

4.4Ablation study

Ablation study of ASS-Sub-Agent. To demonstrate the impact of ASS-Sub-Agent, we compare the qualitative results of our method with and without ASS-Sub-Agent in Fig. 10. The second row shows our results with ASS-Sub-agent. Apparently, under the same number of training epochs, the training results with ASS-Sub-Agent are better than those without it. Because the sub-agents cannot share the training process without ASS-Sub-Agent and each sub-agent is trained for only 1/n of the total epochs.

Figure 10.

Qualitative results of the ablation study of ASS-Sub-Agent. The first row shows the target images, the second row and the third row show the mesh results of ours with and without ASS-Sub-Agent separately.

Ablation study of hyperparameters. We refer to the baseline to set the hyperparameters, which can also ensure fairness. To verify the appropriateness of the selected hyperparameters, we performed an ablation study. As shown in the Table 6, we test the values that are smaller and greater than the selected value respectively and obtain relative results. It can be seen that the results with other values are worse than the results with selected values, indicating that the selected values of hyperparameters are reasonable. Specifically, for the discounted factor γ, when γ= 0.5, the agent is slightly short-sighted because too much attention is given to immediate rewards. When γ= 1, The cumulative reward is too affected by uncertainties in the future, which easily leads to nonconvergence. For the selection of loop number, setting it to 10 can achieve relatively high resolution. Setting the greedy probability ϵ to 0.98 achieves a balance between exploration and exploitation. As shown in the Fig. 11, when ϵ= 1, the initial results are poor because no exploration. The initial result is better when ϵ= 0.5, but it is difficult to reach the optimal result.

Figure 11.

The influence of the greedy probability. The charts show the change of cumulative reward with the training epoch using different probabilities.