Robot collision detector built on pygame with CNN Pytorch implementation
This project involves training a robot using a feed-forward neural network for action conditioning. The training process includes data loading, model training, and result visualization.
Data_Loaders
: Provides data loaders for training and testing the robot model.Networks
: Contains the implementation of the action-conditioned feed-forward neural network.torch
: PyTorch library for neural network training and evaluation.torch.nn
: Provides various neural network modules and loss functions.matplotlib.pyplot
: Plotting library for visualizations.torch.optim
: Optimization algorithms for updating model parameters.numpy
: Numerical computing library.
The train_model
function trains the robot model. It takes the number of epochs as an input parameter and performs the following steps:
- Initialize data loaders, model, optimizer, and loss function.
- Iterate over the specified number of epochs and perform model training.
- Calculate the training loss for each epoch and store it in
training_losses
. - Print the test loss and accuracy for every 10th epoch.
- Save the trained model to a file.
- Plot the training and test loss curves.
The train_model
function is called from the __main__
block, where the number of epochs is specified. This block serves as the entry point for executing the training process.
To run the project, ensure that all dependencies are installed and execute the script. Adjust the number of epochs as needed.