dnnsdunca commited on
Commit
6a31365
1 Parent(s): 603cc72

Create models/perception_agent.py

Browse files
Files changed (1) hide show
  1. models/perception_agent.py +23 -0
models/perception_agent.py ADDED
@@ -0,0 +1,23 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+
4
+ class PerceptionAgent(nn.Module):
5
+ def __init__(self, config):
6
+ super(PerceptionAgent, self).__init__()
7
+ self.cnn_layers = nn.Sequential(
8
+ nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),
9
+ nn.ReLU(),
10
+ nn.MaxPool2d(kernel_size=2, stride=2),
11
+ # Additional layers can be defined based on config
12
+ )
13
+ self.fc_layers = nn.Sequential(
14
+ nn.Linear(16 * 32 * 32, 256),
15
+ nn.ReLU(),
16
+ nn.Linear(256, config["perception_output_size"])
17
+ )
18
+
19
+ def forward(self, x):
20
+ x = self.cnn_layers(x)
21
+ x = x.view(x.size(0), -1)
22
+ x = self.fc_layers(x)
23
+ return x