# How to Train the Paint Agent in Learning-to-Paint: Key Source Files and Architecture

> Learn to train the paint agent in hzwer/iccv2019-learningtopaint. Discover key source files and the training architecture involving DDPG networks, environment simulators, and training scripts.

- Repository: [hzwer/iccv2019-learningtopaint](https://github.com/hzwer/iccv2019-learningtopaint)
- Tags: how-to-guide
- Published: 2026-03-03

---

**Training the paint agent in the hzwer/iccv2019-learningtopaint repository requires orchestrating three architectural layers: training scripts that manage hyperparameters and TensorBoard logging, environment simulators that render brush strokes onto a canvas, and Deep Deterministic Policy Gradient (DDPG) networks with auxiliary GAN rewards that learn sequential decision-making.**

The ICCV 2019 Learning-to-Paint project implements a reinforcement learning agent that reproduces images through sequential brush strokes. Understanding which source files control the training pipeline allows you to modify hyperparameters, replace network backbones, or adjust the perceptual reward function without breaking the system.

## Training Orchestration: Entry Points for Paint Agent Training

The training process begins in the orchestration layer, which parses command-line arguments, initializes logging, and manages the global training loop.

**Primary training scripts** are located at [`baseline/train.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/train.py) and [`baseline_modelfree/train.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline_modelfree/train.py). These files handle argument parsing (lines 78-98) to create a configuration namespace `args`, instantiate a TensorBoard logger via `TensorBoard('../train_log/{}'.format(exp))`, and launch the vectorized environment through `fastenv` from [`baseline/DRL/multi.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/multi.py).

The core training logic resides in the `train()` function (lines 30-75), which implements the standard RL loop:

1. Selects actions via `agent.select_action`
2. Steps the environment with `env.step`
3. Stores transitions in replay memory using `agent.observe`
4. Periodically updates policy parameters through `agent.update_policy`
5. Validates performance using the `Evaluator` class and logs metrics to TensorBoard

## Environment Simulation and Stroke Rendering

Before the agent can learn, the environment must simulate the painting canvas and render parameterized brush strokes.

**Environment definitions** live in [`baseline/env.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/env.py) and [`baseline_modelfree/env.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline_modelfree/env.py). These modules supply the state representation consisting of the current canvas, target image patch, step counter, and coordinate map. The environment's `step()` method accepts brush parameters from the actor network and returns the next canvas state along with the reward signal.

**Stroke rendering** is handled by [`baseline/Renderer/model.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/Renderer/model.py) and [`baseline/Renderer/stroke_gen.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/Renderer/stroke_gen.py). The `FCN` class and `Decoder` network decode the 13-dimensional stroke parameters (10 shape parameters plus 3 color channels) into pixel-space brush strokes. The `decode` function in [`baseline/DRL/ddpg.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/ddpg.py) blends these strokes onto the canvas using differentiable rendering, enabling gradient flow back to the policy network.

## Deep RL Agent: DDPG, Actor-Critic Networks, and GAN Rewards

The intelligence layer implements the DDPG algorithm with an auxiliary Wasserstein GAN that provides perceptual rewards.

**Agent core** is implemented in [`baseline/DRL/ddpg.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/ddpg.py). This file defines the `DDPG` class (instantiated at lines 106-111 in [`train.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/train.py)) that orchestrates the actor-critic architecture. During training, the `update_policy()` method computes standard Q-learning targets while also calling `update_gan()` to compute perceptual rewards that encourage strokes improving visual similarity to the target.

**Network architectures** are defined in [`baseline/DRL/actor.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/actor.py) and [`baseline/DRL/critic.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/critic.py). Both use ResNet backbones: the actor network predicts continuous brush parameters (position, size, color), while the critic network estimates Q-values for state-action pairs.

**Supporting modules** include:

- [`baseline/DRL/rpm.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/rpm.py): Implements the replay memory buffer for storing and sampling transitions
- [`baseline/DRL/wgan.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/wgan.py): Contains the stroke-level Wasserstein GAN that supplies auxiliary perceptual rewards
- [`baseline/DRL/evaluator.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/evaluator.py): Provides validation utilities for testing policy performance on held-out images
- [`baseline/utils/tensorboard.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/utils/tensorboard.py) and [`baseline/utils/util.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/utils/util.py): Handle logging infrastructure and utility functions

## Running the Training Pipeline

Execute a full training session using the baseline implementation with the following command:

```bash
python baseline/train.py \
    --env_batch 96 \
    --train_times 2000000 \
    --max_step 40 \
    --batch_size 96 \
    --rmsize 800 \
    --noise_factor 0.1 \
    --validate_interval 50 \
    --output ./model

```

**Key parameters explained:**

- `--env_batch`: Number of parallel environments running simultaneously (96)
- `--train_times`: Total number of training steps (2,000,000)
- `--max_step`: Maximum strokes per episode (40)
- `--rmsize`: Replay memory size (800)
- `--noise_factor`: Exploration noise magnitude in parameter space (0.1)
- `--validate_interval`: Validation frequency in steps (50)

## Debugging and Inspecting Agent Components

To verify the network architecture before training:

```python
from baseline.DRL.ddpg import DDPG

agent = DDPG(batch_size=96, env_batch=1, max_step=40)
print(agent.actor)    # ResNet actor architecture

print(agent.critic)   # ResNet critic architecture

```

For debugging stroke rendering without the full training loop:

```python
from baseline.Renderer.model import FCN
from baseline.DRL.ddpg import decode
import torch

decoder = FCN()
decoder.load_state_dict(torch.load('../renderer.pkl'))

# Generate random stroke parameters: 10 shape + 3 color

stroke_params = torch.randn(1, 13)
canvas = torch.zeros(1, 3, 128, 128)

# Apply stroke to blank canvas

canvas = decode(stroke_params, canvas)

```

## Summary

- **Training orchestration** is controlled by [`baseline/train.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/train.py) and [`baseline_modelfree/train.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline_modelfree/train.py), which manage hyperparameters, TensorBoard logging, and the global training loop.
- **Environment simulation** occurs in [`baseline/env.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/env.py) and [`baseline_modelfree/env.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline_modelfree/env.py), providing canvas state and reward signals, while [`baseline/Renderer/model.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/Renderer/model.py) and [`stroke_gen.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/stroke_gen.py) handle differentiable stroke rendering.
- **Agent intelligence** is implemented in [`baseline/DRL/ddpg.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/ddpg.py), utilizing actor networks ([`actor.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/actor.py)), critic networks ([`critic.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/critic.py)), replay memory ([`rpm.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/rpm.py)), and a perceptual GAN reward ([`wgan.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/wgan.py)).
- **Validation and utilities** are provided by [`baseline/DRL/evaluator.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/evaluator.py) and helper modules in `baseline/utils/`.

## Frequently Asked Questions

### Where is the main training loop defined for the paint agent?

The main training loop is defined in the `train()` function within [`baseline/train.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/train.py) (lines 30-75). This function orchestrates the interaction between the DDPG agent, vectorized environments from [`baseline/DRL/multi.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/multi.py), and the replay buffer, while periodically triggering validation through the `Evaluator` class and logging metrics to TensorBoard.

### What is the difference between the baseline and baseline_modelfree training scripts?

Both [`baseline/train.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/train.py) and [`baseline_modelfree/train.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline_modelfree/train.py) implement the same DDPG training logic, but they differ in their environment implementations. The baseline version uses [`baseline/env.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/env.py) which may include model-based components or specific stroke parameterizations, while the model-free variant in `baseline_modelfree/` uses a simplified environment definition in [`baseline_modelfree/env.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline_modelfree/env.py) that relies purely on learned policies without explicit stroke models.

### How does the GAN reward function improve paint agent training?

The GAN reward, implemented in [`baseline/DRL/wgan.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/wgan.py), provides a perceptual loss signal computed by a Wasserstein GAN during the `update_gan()` call within the DDPG update cycle. This auxiliary reward supplements the standard Q-learning target by evaluating whether individual strokes improve the perceptual similarity between the current canvas and the target image, encouraging more visually coherent painting strategies.

### Which files should I modify to change the neural network architecture?

To modify the paint agent's neural architecture, edit [`baseline/DRL/actor.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/actor.py) for the policy network (which determines brush stroke parameters) and [`baseline/DRL/critic.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/critic.py) for the value network. Both files define ResNet-based architectures that map environment states to actions or Q-values. Changes to these files automatically propagate through [`baseline/DRL/ddpg.py`](https://github.com/hzwer/iccv2019-learningtopaint/blob/main/baseline/DRL/ddpg.py), which instantiates these networks during agent initialization.