environment.py

import logging
import numpy as np
import pandas as pd
import torch
from typing import Dict, List, Optional, Tuple
from abc import ABC, abstractmethod
from enum import Enum
from dataclasses import dataclass
from scipy.optimize import minimize
from scipy.spatial import distance
from scipy.stats import norm
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from sklearn.utils import shuffle
from torch.utils.data import Dataset, DataLoader
from torch.optim import Adam
from torch.nn import Module, Linear, ReLU, Sigmoid, BCELoss
from torch.nn.functional import softmax
from torch.optim.lr_scheduler import StepLR
from torch.utils.tensorboard import SummaryWriter

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Constants and configuration
class Config(Enum):
    TRAINING_EPOCHS = 100
    BATCH_SIZE = 32
    LEARNING_RATE = 0.001
    MARGIN = 0.1
    THRESHOLD = 0.5
    NUM_FEATURES = 10
    NUM_CLASSES = 2

@dataclass
class EnvironmentConfig:
    training_epochs: int = Config.TRAINING_EPOCHS.value
    batch_size: int = Config.BATCH_SIZE.value
    learning_rate: float = Config.LEARNING_RATE.value
    margin: float = Config.MARGIN.value
    threshold: float = Config.THRESHOLD.value
    num_features: int = Config.NUM_FEATURES.value
    num_classes: int = Config.NUM_CLASSES.value

class Environment(ABC):
    def __init__(self, config: EnvironmentConfig):
        self.config = config
        self.model = None
        self.data = None
        self.loader = None
        self.writer = None

    @abstractmethod
    def setup(self):
        pass

    @abstractmethod
    def interact(self):
        pass

    def train(self):
        if self.model is None:
            raise ValueError("Model not initialized")
        if self.data is None:
            raise ValueError("Data not loaded")
        if self.loader is None:
            raise ValueError("Data loader not created")

        self.model.train()
        for epoch in range(self.config.training_epochs):
            for batch in self.loader:
                inputs, labels = batch
                self.model.zero_grad()
                outputs = self.model(inputs)
                loss = BCELoss()(outputs, labels)
                loss.backward()
                self.model.step()
            self.writer.add_scalar("Loss", loss.item(), epoch)

    def evaluate(self):
        if self.model is None:
            raise ValueError("Model not initialized")
        if self.data is None:
            raise ValueError("Data not loaded")
        if self.loader is None:
            raise ValueError("Data loader not created")

        self.model.eval()
        total_loss = 0
        with torch.no_grad():
            for batch in self.loader:
                inputs, labels = batch
                outputs = self.model(inputs)
                loss = BCELoss()(outputs, labels)
                total_loss += loss.item()
        return total_loss / len(self.loader)

class DiscreteTimeEnvironment(Environment):
    def __init__(self, config: EnvironmentConfig):
        super().__init__(config)
        self.setup()

    def setup(self):
        self.model = torch.nn.Sequential(
            Linear(self.config.num_features, 128),
            ReLU(),
            Linear(128, 128),
            ReLU(),
            Linear(128, self.config.num_classes),
            Sigmoid()
        )
        self.model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
        self.data = torch.randn(self.config.num_features, self.config.num_classes)
        self.loader = DataLoader(self.data, batch_size=self.config.batch_size, shuffle=True)

    def interact(self):
        inputs = torch.randn(self.config.num_features, 1)
        outputs = self.model(inputs)
        return outputs.item()

class ContinuousTimeEnvironment(Environment):
    def __init__(self, config: EnvironmentConfig):
        super().__init__(config)
        self.setup()

    def setup(self):
        self.model = torch.nn.Sequential(
            Linear(self.config.num_features, 128),
            ReLU(),
            Linear(128, 128),
            ReLU(),
            Linear(128, self.config.num_classes),
            Sigmoid()
        )
        self.model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
        self.data = torch.randn(self.config.num_features, self.config.num_classes)
        self.loader = DataLoader(self.data, batch_size=self.config.batch_size, shuffle=True)

    def interact(self):
        inputs = torch.randn(self.config.num_features, 1)
        outputs = self.model(inputs)
        return outputs.item()

class FlowTheoryEnvironment(Environment):
    def __init__(self, config: EnvironmentConfig):
        super().__init__(config)
        self.setup()

    def setup(self):
        self.model = torch.nn.Sequential(
            Linear(self.config.num_features, 128),
            ReLU(),
            Linear(128, 128),
            ReLU(),
            Linear(128, self.config.num_classes),
            Sigmoid()
        )
        self.model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
        self.data = torch.randn(self.config.num_features, self.config.num_classes)
        self.loader = DataLoader(self.data, batch_size=self.config.batch_size, shuffle=True)

    def interact(self):
        inputs = torch.randn(self.config.num_features, 1)
        outputs = self.model(inputs)
        return outputs.item()

class VelocityThresholdEnvironment(Environment):
    def __init__(self, config: EnvironmentConfig):
        super().__init__(config)
        self.setup()

    def setup(self):
        self.model = torch.nn.Sequential(
            Linear(self.config.num_features, 128),
            ReLU(),
            Linear(128, 128),
            ReLU(),
            Linear(128, self.config.num_classes),
            Sigmoid()
        )
        self.model.to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
        self.data = torch.randn(self.config.num_features, self.config.num_classes)
        self.loader = DataLoader(self.data, batch_size=self.config.batch_size, shuffle=True)

    def interact(self):
        inputs = torch.randn(self.config.num_features, 1)
        outputs = self.model(inputs)
        return outputs.item()

if __name__ == "__main__":
    config = EnvironmentConfig()
    env = DiscreteTimeEnvironment(config)
    env.train()
    env.evaluate()