RL-GAN-Net: A Reinforcement Learning Agent Controlled GAN Network for Real-Time Point Cloud Shape Completion

In plain words

A 3D scanner often misses parts of an object. Some areas are hidden from the sensor, some surfaces reflect badly, and the result is a cloud of points with holes. This paper fills those holes. It takes the partial point cloud and returns a complete shape with the missing parts filled in.

The method has three learned pieces. A first network compresses a shape into a short list of numbers, and a second network turns that list back into points. A third network, a GAN, can produce plausible short lists from a random seed. The hard part is choosing the right seed for the partial scan. Earlier work searched for it by optimisation, which has to be repeated for each batch. Here a small reinforcement-learning agent learns to pick the seed directly from the scan, so the completion takes one fast pass through the networks and runs in real time.

For scans with few holes, a plain autoencoder is already good. The paper therefore also offers a hybrid version that uses the GAN's discriminator to choose between the GAN result and the autoencoder result.

TL;DR

RL-GAN-Net trains an autoencoder, then a latent GAN on its codes, then a DDPG actor-critic agent that outputs the one-dimensional GAN seed for each partial input §3.1–§3.3. At inference, a partial cloud is encoded, the agent picks the seed, and the GAN output is decoded alg. 1. Use it as a reference for real-time completion of 2048-point ShapeNet-style shapes with 20–70% missing points. All numbers are author-reported. The official code omits the hybrid switch and departs from the paper's reward and Chamfer definitions, so read [Paper vs code] before reimplementing.

Problem

Acquired 3D data (laser scans, stereo, RGB-D) is a point cloud with large missing regions, caused by limited viewpoints, occlusion, sensor resolution, texture-less surfaces or specular materials §1. Point-cloud shape completion is a data-driven fix, and point clouds suit it because voxel grids grow with the cube of resolution §1.

Two families of prior work are relevant. Autoencoders decode a complete shape from an incomplete encoding, but quality falls as the missing fraction grows §3.1. Latent-GAN completion searches for the GAN input that best explains the partial cloud, by an optimisation that is run for each batch of inputs §3.2 §4.1. The paper states the reported cost of that baseline as 324 seconds for a batch of 50 shapes §4.1, and the conclusion again calls it "a minimum of 324 seconds" §5.

The paper's assumption is that the GAN should be trained on the latent codes, not on raw points, because training on latents is more stable §3.2. It also assumes that the agent can learn the nonlinear map from noisy codes to good GAN seeds [abstract]. For small missing fractions the autoencoder alone can recover the shape well, which motivates the hybrid §3.4.

Method

Autoencoder (AE) §3.1. The encoder maps a 2048-point cloud to a global feature vector (GFV) of 128 values. The decoder maps a GFV back to 2048 points [supp. AE Details]. Training minimises the symmetric Chamfer distance, given as d_{CH} in eq. 1. The paper prefers Chamfer over the Earth Mover's distance for efficiency §3.1.

d_{CH}(P_1, P_2) = \sum_{a \in P_1} \min_{b \in P_2} \lVert a-b \rVert_2^2 + \sum_{b \in P_2} \min_{a \in P_1} \lVert a-b \rVert_2^2 \qquad [\text{eq. 1}]

Latent GAN, l-GAN §3.2. A GAN is trained on the GFVs of complete shapes, with the AE encoder producing those GFVs §3.2. The generator G maps a noise seed z to a clean GFV, and the decoder turns that GFV into points. The GAN uses a WGAN-GP adversarial loss, with a self-attention generator and discriminator and spectral normalisation [supp. l-GAN Details].

Reward terms §3.3 eq. 2–eq. 5. For a partial input P_{in}, the paper defines three losses. The Chamfer loss compares P_{in} with the decoded generated shape eq. 2. The GFV loss is the squared distance between G(z) and the encoding E(P_{in}) eq. 3. The discriminator loss is -D(G(z)) eq. 4. The reward is the weighted sum of negated losses eq. 5.

L_{CH} = d_{CH}\big(P_{in},\, E^{-1}(G(z))\big) \qquad [\text{eq. 2}]

L_{GFV} = \lVert G(z) - E(P_{in}) \rVert_2^2 \qquad [\text{eq. 3}]

L_{D} = -D(G(z)) \qquad [\text{eq. 4}]

r = w_{CH}\, r_{CH} + w_{GFV}\, r_{GFV} + w_{D}\, r_{D}, \quad r_{\bullet} = -L_{\bullet} \qquad [\text{eq. 5}]

The weights are w_{CH}=100, w_{GFV}=10, w_D=0.01 [supp. RL Agent Details]. The paper states they were chosen so that no term dominates, using the typical loss range over several trials [supp. RL Agent Details].

RL agent §3.3 alg. 1. The environment is the frozen encoder, generator, discriminator and decoder. The state is the noisy GFV of P_{in} (128 values). The action is the one-dimensional seed z [supp. RL Agent Details]. Each episode is one step: the agent emits one z, and the episode ends [supp. RL Agent Details]. The agent is DDPG §3.3: a deterministic actor \mu(s \mid \theta^\mu), a critic Q(s,a) trained with the Bellman target, and the deterministic policy gradient eq. 6. Training is off-policy, with a replay buffer and batches of 100 [supp. RL Agent Details].

\nabla_{\theta^\mu} J \approx \mathbb{E}_{s}\!\left[ \nabla_a Q(s, a \mid \theta^Q)\big|_{a=\mu(s)}\, \nabla_{\theta^\mu} \mu(s \mid \theta^\mu) \right] \qquad [\text{eq. 6}]

Hybrid §3.4. The paper proposes using the GAN discriminator to choose, per input, between the vanilla RL-GAN-Net output and the AE output §3.4 fig. 2. This step is described in text only; the repository has no code for it (see [Paper vs code]).

The training procedure is alg. 1. The AE and GAN are pretrained first. The agent then runs a loop of: observe state, select action, run the environment to get the reward, and store the transition in the replay buffer. The actor and critic are trained from the buffer alg. 1.

# Training the RL agent  [alg. 1]
# Input:  pretrained encoder E, decoder Dec, generator G, discriminator D
# Output: trained DDPG actor and critic

def train_agent(partial_clouds, E, Dec, G, D):
    # Create DDPG actor, critic and empty replay buffer  [§3.3] [alg. 1]
    actor, critic, replay = create_ddpg()
    for step in range(max_steps):
        # Sample an incomplete point cloud from the training data  [alg. 1]
        p_in = next(partial_clouds)                   # (2048, 3)
        # Encode it to a noisy GFV; this is the state  [§3.3] [alg. 1]
        state = E(p_in)                               # (128,)
        # Choose the GAN seed, with exploration noise  [§3.3] [supp. RL]
        z = actor(state) + exploration_noise()        # (1,)
        # Generate a clean GFV from the seed  [§3.2] [alg. 1]
        gfv_clean = G(z)                              # (128,)
        # Losses on the generated shape  [eq. 2] [eq. 3] [eq. 4]
        loss_ch = chamfer(Dec(gfv_clean), p_in)       # scalar
        loss_gfv = squared_l2(gfv_clean, state)       # scalar
        d_score = D(gfv_clean)                        # scalar
        # Reward: sum of weighted negated losses, with r_D = D(G(z))  [eq. 5]
        reward = (-100.0 * loss_ch
                  - 10.0 * loss_gfv
                  + 0.01 * d_score)                   # scalar
        # Store the transition, then update DDPG from a batch of 100  [alg. 1]
        replay.add(state, z, reward, gfv_clean)       # next state: see Pitfalls
        critic, actor = ddpg_update(replay, batch=100)  # [supp. RL Agent]
    return actor, critic

Inference, following alg. 1 "Final output":

# Completing one partial point cloud  [alg. 1] [§3.3]
# Input:  partial point cloud p_in (2048, 3)
# Output: completed point cloud (2048, 3)

def complete(p_in, E, Dec, G, actor):
    # Encode the partial input to a noisy GFV  [alg. 1]
    state = E(p_in)                               # (128,)
    # Pick the seed; no exploration noise at test time  [alg. 1]
    z = actor(state)                              # (1,)
    # Generate a clean GFV and decode it to points  [alg. 1]
    gfv_clean = G(z)                              # (128,)
    return Dec(gfv_clean)                         # (2048, 3)

Hyperparameters

Values in the paper's own tables and text are cited to the paper. Values taken from the released code are marked and cited to the code. Where a value is only in the code and not in the paper, the paper's value is "not stated".

Results

Author-reported, on ShapeNet-style data with the four categories above. Baselines are the AE alone §4.1 [supp. Classifier Details], vanilla PointNet on raw partial input [supp. Classifier Details], and the optimisation method of Gurumurthy et al. for run time §4.1. The text gives no exact Chamfer values for the shape-completion comparison. Those are only in a plot fig. 5.

The classifier is PointNet trained on complete shapes, with 99.36% accuracy on complete clouds [supp. Classifier Details]. The paper states the hybrid is better than vanilla on Chamfer and about equal on classification §4.1 §4.2. A reward ablation shows that Chamfer alone fails to find the right category, that GFV loss alone finds the category but not the detail, and that the discriminator loss alone gives a realistic but unrelated shape §4.3 fig. 8. The ablation has no table of numbers.

Paper vs code

Official repository: https://github.com/iSarmad/RL-GAN-Net (MIT). Its README calls it the "Official Repository" of this paper and links the arXiv id README.md. The LICENSE names Muhammad Sarmad, the paper's first author LICENSE. Commit f0de736 (shallow clone).

The paper and code differ in these places. Each is checked against the code, and the code is taken as what was actually run.

• Reward has a fourth term. The paper's reward eq. 5 has three terms. The code adds -0.1\,(\lVert z \rVert_2 - 1)^2 to the reward, which is a norm penalty on the seed trainRL.py trainRL.py lossess.py. The paper's [supp. RL Agent Details] does not mention it.
• Chamfer target is the AE reconstruction, not the input, in training. The paper defines L_{CH} against P_{in} eq. 2. Training computes it against the AE decoding of the noisy GFV, E^{-1}(E(P_{in})), and the line that uses P_{in} is commented out trainRL.py. The test script uses P_{in} testRL.py.
• Chamfer form differs. The paper sums squared distances eq. 1. The code takes per-point Euclidean distances (square root of the sum of squares) to nearest neighbours and averages them over points and batch, rather than summing lossess.py lossess.py. The averaging changes the scale relative to the GFV term and to the weights.
• GFV loss form differs. The paper uses the squared L2 norm eq. 3. The code uses elementwise MSE, averaging over the 128 dimensions trainRL.py lossess.py. This is smaller by a factor of 128 than the paper's definition, before the weight 10 is applied.
• Every transition is stored as terminal. The discount term is switched off in practice. The environment always returns done = True trainRL.py, and the episode-step counter is reset before each store, so done_bool is 1 for almost all transitions trainRL.py. The paper gives \gamma=0.9 [supp. RL Agent Details]; the code default is 0.99 RL_params.py. As written, \gamma has no effect on the targets.
• Algorithm 1 has the warm-up conditions reversed. alg. 1 lines "t_{steps} > t_{StartTime}: Random Action" and "t_{steps} < t_{StartTime}: Use actor" are the opposite of the code, which uses random actions for the first steps and the actor afterwards trainRL.py. The paper gives no warm-up length; the code uses 1e4 RL_params.py.
• Action bound. The paper says the action space is "kept to unity" [supp. RL Agent Details]. The code clips actions to \pm 10 and warms up with uniform samples in [-10, 10] RL_params.py trainRL.py.
• TD3 parameters listed for DDPG. The paper's table lists target-policy noise 0.2, noise clip 0.5 and delayed policy update every 2 steps [supp. RL Agent Details]. These are TD3 terms. The default policy is DDPG RL_params.py trainRL.py, and DDPG.py has no target-policy noise and no delayed update.
• Critic width. The paper lists critic layers 400, 432, 300, 300 [supp. RL Agent Details]. The code concatenates the action after the first layer, so the second layer is 400 + z_{dim} wide, then 300, then 300, then 1 DDPG.py. The 432 would match a 32-dim action, a z size the paper says it tested but kept at 1.
• Discriminator self-attention. The paper's discriminator table has one self-attention block at 256×4×4 [supp. l-GAN Details]. The code instantiates two, with 256 and 512 channels self_D_net.py.
• Hybrid selection is absent. The paper describes the discriminator switch §3.4 and the pipeline figure fig. 2. No code in the repository implements it. testRL.py only produces the agent's output testRL.py. A grep for "hybrid" and "switch" in the Python sources found no selection logic.
• Checkpoints are not included. The scripts load the AE, GAN and actor from the author's local paths RL_params.py. The README lists the stages: download the dataset, preprocess, train AE, generate GFVs, train GAN, train RL README.md. Preprocessing scripts are MATLAB (processData.m, processData2.m).
• Test-time reward is different. testRL.py sums the losses unweighted in its reward and uses P_{in} as the Chamfer target testRL.py testRL.py. This only affects the reported test loss, not training.

Several items the paper gives are consistent with the code. These are the DDPG \tau=0.005, exploration noise 0.1, replay batch 100, 10^6 steps, the GAN optimiser and learning rate, the GAN batch size of 50, \lambda_{gp}=10, five discriminator steps, and the reward weights 100, 10, 0.01 [code: RL_params.py#L104-L109] parameter.py [code: trainRL.py#L436].

Dependencies the paper does not mention: the Chamfer code uses FAISS GPU indices lossess.py lossess.py, and the RL code calls .cuda() directly trainRL.py. The code targets Python 3.6 (cached .pyc files) with the packages in Requirements_Pip.txt Requirements_Pip.txt.

Pitfalls

• Pick one reward and one Chamfer definition on purpose. The paper's eq. 1–eq. 5 and the code disagree on the Chamfer target, the Chamfer form (sum of squares vs averaged distances), the GFV form (squared norm vs elementwise MSE), and the reward terms. The relative size of the terms sets what the agent learns. The paper's weights (100, 10, 0.01) were chosen for the paper's definitions, so they cannot be copied onto a different loss scale without a check [supp. RL Agent Details].
• Next-state definition. The paper's Algorithm 1 stores (s_t, a_t, r_t, s_{t+1}) but does not define s_{t+1} alg. 1. The code stores the generated clean GFV trainRL.py. With every transition terminal trainRL.py, this value is not used in the targets. Any non-terminal implementation needs a choice here, and the paper does not give it.
• Seed range. The GAN generator's seed is sampled as a standard normal in training trainer.py, while the agent's actions are clipped to \pm 10 RL_params.py. Values well outside the training seed distribution are not checked by the paper. Whether the agent stays in the useful range depends on the actor's output.
• Single-step episodes. Each episode is one action [supp. RL Agent Details], and the code's reset makes each transition terminal. The DDPG update is then a contextual-bandit regression on the reward, not a sequential RL problem. The paper's "episodic" wording and its DDPG framing overstate the sequential structure.
• Unstated optimiser for the AE. The paper gives no learning rate, optimiser or epoch count for the AE [supp. AE Details]. The code defaults are batch 24, lr 0.001 and 400 epochs main.py. A pretrained-checkpoint path in the code name suggests other values; that path is local to the author and is not documented, so do not rely on it.
• Data generation and jitter. Partial inputs are made by removing points inside a radius around a random seed point, with the radius set per shape to reach the target fraction §4. Jitter is added to training inputs in the code (Gaussian, \sigma = 0.01, clipped at 0.05) [code: trainRL.py#L97]; the paper mentions this jitter only for the robustness test [supp. Robustness Results]. Whether jitter was part of the reported main results is not stated.
• Chamfer is not a sufficient metric. The paper notes that for 20% missing the input point cloud has a lower Chamfer distance than all completions, while the completions are visually more plausible §4.1. Judge completions with classification or visual checks as well.
• Hidden cost of the optimisation baseline. The 324 s figure is the paper's quote for the baseline on its batch of 50 §4.1 §5. It is not a measurement on this paper's hardware. The paper's own timing is for the forward pass on a single Titan Xp §4 tab. 1.
• Figures carry some numbers. The Chamfer curves for AE and both variants are only plotted fig. 5. A reimplementation cannot use text values for those comparisons.
• Table column order. The timing table lists the ratios as 20, 40, 30, 50, 70 tab. 1. Read them by ratio, not by position.

Lineage

• Builds on: Achlioptas et al., autoencoder and latent-GAN for point clouds, which the paper's AE and latent GAN follow §3.1 §3.2. Gurumurthy et al., latent-GAN completion by optimising the GAN seed for each input; this paper replaces that search with a trained agent §2 §3.2. Lillicrap et al., DDPG §3.3. The self-attention GAN and WGAN-GP cited for the GAN [supp. l-GAN Details]. Dai et al., voxel-based completion, used as a comparison §2 §4.1.
• Compared against: the AE-only completion §4.1; Gurumurthy et al. for run time §4.1; Dai et al. on the released robustness data [supp. Robustness Results]; PointNet classification on raw and completed inputs §4.2.
• Superseded by: not stated in the paper, and no later wiki page covers it.