# Training

This document describes the DQN training workflow implemented in `train.py`.

The training script initializes the custom 3D path-planning environment, creates the DQN policy and target networks, trains the policy using experience replay, and saves the trained model checkpoints.

## Overview

The training workflow uses a standard Deep Q-Network setup:

1. Create the environment.
2. Reset the environment and process the state into a PyTorch tensor.
3. Select actions using an epsilon-greedy policy.
4. Store transitions in replay memory.
5. Sample batches from replay memory.
6. Optimize the policy network.
7. Soft-update the target network.
8. Track episode success and episode length.
9. Save the trained policy and target networks.

The script can be run locally or in Google Colab. When running in Colab, the script mounts Google Drive and uses the Colab/Drive path for output storage.

## Main Functions

## `select_action`

```python
select_action(state, steps_done, policy_net, env, device)
```

Selects an action using an epsilon-greedy strategy.

The function balances exploration and exploitation using the epsilon schedule defined in `dqn.py`:

- `EPS_START`
- `EPS_END`
- `EPS_DECAY`

During exploitation, the policy network selects the action with the highest predicted Q-value. During exploration, a random action is sampled from the environment action space.

## `optimize_model`

```python
optimize_model(memory, policy_net, target_net, optimizer, device)
```

Optimizes the policy network using a batch of transitions sampled from replay memory.

The optimization process:

1. Samples a batch from `ReplayMemory`.
2. Separates state, action, next-state, and reward values.
3. Computes Q-values for the selected actions.
4. Computes target Q-values using the target network.
5. Calculates the Huber loss.
6. Backpropagates the loss.
7. Clips gradients.
8. Updates the policy network.

If the replay memory contains fewer transitions than `BATCH_SIZE`, the function returns without optimizing.

## `plot_steps`

```python
plot_steps(episode_steps, episode_success, is_ipython, show_result=False)
```

Plots training progress by showing the number of steps taken per episode.

Successful and unsuccessful episodes are marked separately, making it easier to inspect whether the agent is learning to reach the goal more efficiently over time.

## `main`

```python
main()
```

Runs the full training loop.

The `main()` function:

- selects the compute device
- creates the `World` environment
- initializes the policy and target networks
- creates the optimizer
- creates the replay memory
- assigns random goals during training
- runs the episode loop
- updates the DQN policy
- soft-updates the target network
- plots training progress
- saves the trained model checkpoints

## Environment Setup

The training script creates the environment using the default goal and start position imported from `dqn.py`:

```python
env = World(goal, start_position, verbosity)
```

The default training state removes the `"prev_position"` key before converting the state into a tensor:

```python
avoid = ["prev_position"]
n_observations = len(env.reset().process(device, avoid))
```

This processed state size is used as the input size for the DQN model.

## Model Setup

The script uses `DQN3` as the default model architecture:

```python
policy_net = DQN3(n_observations, n_actions, 512).to(device)
target_net = DQN3(n_observations, n_actions, 512).to(device)
target_net.load_state_dict(policy_net.state_dict())
```

The policy network is trained directly. The target network is used to compute target Q-values and is updated gradually using a soft update rule.

## Training Episodes

The number of training episodes depends on the available device:

| Device | Episodes |
|---|---|
| CUDA GPU available | `1500` |
| CPU only | `1` |

The CPU setting is mainly useful for checking that the script runs. Meaningful training is expected to use a GPU.

At the start of each episode, a random goal is assigned:

```python
env.assign_random_goal(distance_val=range(250, 1001))
```

This trains the agent across a range of nearby goal positions rather than a single fixed target.

## Target Network Updates

The target network is updated using a soft update rule:

```text
target = TAU * policy + (1 - TAU) * target
```

This helps stabilize DQN training by preventing the target values from changing too abruptly.

## Model Outputs

At the end of training, the script saves two checkpoint files:

```text
<file_id>_policy_net.pth
<file_id>_target_net.pth
```

The policy network is the main model used for action selection after training.

The target network is primarily used during training, but it is also saved so the full training state can be inspected or reused.

## Output Directories

When running locally, model checkpoints are saved to:

```text
models/
```

When running in Google Colab, outputs are saved using the Colab/Drive output path.

The environment also uses the logging directory defined in `environment.py`. Runtime logs should be written to the project-level `logs/` directory when running locally from the project root.

## Usage

Run the training script from the project root:

```bash
python -m dqn_path_planning.train
```

If running the file directly instead, make sure the imports and working directory are set correctly for the local project structure.

## Notes

- The script is designed to work both locally and in Google Colab.
- Training on CPU is limited to one episode by default and is mainly useful as a smoke test.
- Training on GPU runs for 1500 episodes by default.
- The default model architecture is `DQN3` with 512 hidden nodes.
- The state passed into the DQN excludes `"prev_position"`.
- Random goals are assigned during training to encourage generalization across nearby goal locations.
- The final policy checkpoint is the main artifact used by the evaluation and path-planning workflows.