Update README.md

README.md

@@ -41,7 +41,7 @@ This is a reinforcement learning model trained to predict real-time sales conver
 - **Framework**: Stable Baselines3 (PPO)
 - **State Representation**: Azure OpenAI embeddings
 - **Action Space**: Continuous (conversion probability 0-1)
-- **Feature Extractor**: Custom CNN layers
+- **Feature Extractor**: Custom Linear layers
 
 ## Quick Start
 
@@ -83,12 +83,12 @@ model_path = hf_hub_download(
 # Check for GPU
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Custom CNN class
-class CustomCNN(BaseFeaturesExtractor):
+# Custom Linear Layer class
+class CustomLN(BaseFeaturesExtractor):
     def __init__(self, observation_space, features_dim: int = 64):
         super().__init__(observation_space, features_dim)
         n_input_channels = observation_space.shape[0]
-        self.cnn = nn.Sequential(
+        self.ln = nn.Sequential(
             nn.Linear(n_input_channels, 512),
             nn.ReLU(),
             nn.Linear(512, 256),
@@ -98,7 +98,7 @@ class CustomCNN(BaseFeaturesExtractor):
         ).to(device)
         
     def forward(self, observations: torch.Tensor) -> torch.Tensor:
-        return self.cnn(observations)
+        return self.ln(observations)
 
 @dataclass
 class SalesAgent:
@@ -112,7 +112,7 @@ class SalesAgent:
         policy_kwargs = dict(
             activation_fn=nn.ReLU,
             net_arch=[dict(pi=[128, 64], vf=[128, 64])],
-            features_extractor_class=CustomCNN,
+            features_extractor_class=CustomLN,
             features_extractor_kwargs=dict(features_dim=64)
         )