hackathonf23-Stacks/removed_hidden_layer.py

import numpy as np
import pandas as pd
import tensorflow as tf
from keras.layers import Dense
from sklearn.preprocessing import MinMaxScaler
import dagshub
import mlflow
import mlflow.keras
import pickle

mlflow.set_tracking_uri('https://dagshub.com/ML-Purdue/hackathonf23-Stacks.mlflow')
dagshub.init(repo_owner='ML-Purdue', repo_name='hackathonf23-Stacks', mlflow=True)

def get_or_create_experiment_id(name):
    exp = mlflow.get_experiment_by_name(name)
    if exp is None:
        exp_id = mlflow.create_experiment(name)
        return exp_id
    return exp.experiment_id
    
class MaxEntIRL:
    def __init__(self, state_dim):
        self.state_dim = state_dim
        self.model = self._create_irl_model()

    def _create_irl_model(self):
        model = tf.keras.Sequential([
            Dense(self.state_dim, input_shape=(self.state_dim,), activation='relu'),
            Dense(4096, activation='relu'),
            Dense(self.state_dim, activation='linear')
        ])
        return model

    def generateHumanTrajectories(self, num_trajectories, trajectory_length):
        human_trajectories = []
        for _ in range(num_trajectories):
            trajectory = []
            state = np.zeros(self.state_dim)
            for _ in range(trajectory_length):
                direction_probabilities = self._generate_direction_probabilities()  # Get direction probabilities
                action_coefficients = np.random.choice([-1, 0, 1], p=direction_probabilities)
                action = action_coefficients * 0.1
                new_state = state + action
                trajectory.append((state, action))
                state = new_state

            human_trajectories.append(trajectory)

        return human_trajectories
    
    def _generate_direction_probabilities(self):
        probabilities = np.random.dirichlet(np.ones(self.state_dim) * 0.1)
        return probabilities

    def loadDataset(self, file_path):
        data = pd.read_csv(file_path)  # Load CSV data
        scaler = MinMaxScaler()
        columns_to_normalize = ['position x [mm]', 'position y [mm]', 'position z (height) [mm]', 'velocity [mm/s]']
        data[columns_to_normalize] = scaler.fit_transform(data[columns_to_normalize])
        return data

    def train_irl_with_dataset(self, data, lr=0.001, epochs=3):
        state_dim = self.state_dim
        optimizer = tf.keras.optimizers.Adam(learning_rate=lr)

        # Extract relevant columns from the loaded dataset
        positions = data[['position x [mm]', 'position y [mm]', 'position z (height) [mm]']].values
        velocities = data['velocity [mm/s]'].values
        mlflow.tensorflow.autolog()
        with mlflow.start_run(experiment_id=get_or_create_experiment_id("Base Reproduction")):
            for epoch in range(epochs):
                total_loss = 0
                state_frequencies = self._calculate_state_frequencies(positions)

                for idx in range(len(positions)):
                    state = positions[idx]
                    velocity = velocities[idx]

                    with tf.GradientTape() as tape:
                        preferences = self.model(state[np.newaxis, :])
                        prob_human = tf.nn.softmax(preferences)

                        # Define losses
                        max_entropy_loss = -tf.reduce_sum(prob_human * tf.math.log(prob_human + 1e-8), axis=1)
                        alignment_loss = -tf.reduce_sum(state_frequencies * tf.math.log(prob_human + 1e-8), axis=1)
                        maxent_irl_objective = max_entropy_loss + alignment_loss

                        # Compute the gradients
                        grads = tape.gradient(maxent_irl_objective, self.model.trainable_variables)
                        optimizer.apply_gradients(zip(grads, self.model.trainable_variables))
                        total_loss += tf.reduce_sum(maxent_irl_objective)  # Accumulate the total loss

                avg_loss = total_loss / len(positions)
                mlflow.log_metric(f"loss_epoch_{epoch}", avg_loss)
                print(f"Epoch {epoch + 1}/{epochs}, MaxEnt IRL Loss: {avg_loss}")

    def train_irl(self, human_trajectories=None, data=None, use_dataset=False, lr=0.001, epochs=3):
        if use_dataset and data is not None:
            # Train using the loaded dataset
            self.train_irl_with_dataset(data, lr=lr, epochs=epochs)
        else:
            # Train using the generative function
            if human_trajectories is None:
                human_trajectories = self.generateHumanTrajectories(num_trajectories, trajectory_length)

            self._train_irl_generative(human_trajectories, lr=lr, epochs=epochs)

    def _train_irl_generative(self, human_trajectories, lr=0.001, epochs=3):
        trajectory_length = len(human_trajectories[0])
        optimizer = tf.keras.optimizers.Adam(learning_rate=lr)

        for epoch in range(epochs):
            total_loss = 0
            state_frequencies = self._calculate_state_frequencies(human_trajectories, trajectory_length)

            for trajectory in human_trajectories:
                for state, _ in trajectory:
                    with tf.GradientTape() as tape:
                        preferences = self.model(state[np.newaxis, :])
                        prob_human = tf.nn.softmax(preferences)

                        # Inside the training loop:
                        max_entropy_loss = -tf.reduce_sum(prob_human * tf.math.log(prob_human + 1e-8), axis=1)
                        alignment_loss = -tf.reduce_sum(state_frequencies * tf.math.log(prob_human + 1e-8), axis=1)
                        maxent_irl_objective = max_entropy_loss + alignment_loss

                    grads = tape.gradient(maxent_irl_objective, self.model.trainable_variables)
                    optimizer.apply_gradients(zip(grads, self.model.trainable_variables))
                    total_loss += maxent_irl_objective

            avg_loss = total_loss / (len(human_trajectories) * trajectory_length)
            mlflow.log_metric(f"loss_epoch_{epoch}", avg_loss)
            autolog_run = mlflow.last_active_run()
            
            print(f"Epoch {epoch + 1}/{epochs}, MaxEnt IRL Loss: {avg_loss}")

    def _calculate_state_frequencies(self, positions):
        state_counts = np.sum(positions, axis=0)
        state_frequencies = state_counts / (len(positions) * self.state_dim)
        return state_frequencies
    def save_model(self, file_path):
        model_config = self.model.get_config()
        with open(file_path, 'wb') as f:
            pickle.dump(model_config, f)

    @classmethod
    def load_model(cls, file_path, state_dim):
        with open(file_path, 'rb') as f:
            model_config = pickle.load(f)
            irl_instance = cls(state_dim)
            irl_instance.model = tf.keras.Sequential.from_config(model_config)
            return irl_instance
    # Indicate test completion status
    
state_dim = 3  # Dimension of the state space
irl = MaxEntIRL(state_dim)
num_trajectories = 100
trajectory_length = 20

# Load the dataset
file_path = '/Users/vinay/Desktop/Computer_Science_Projects/ReScience/hackathonf23-Stacks/data/train.csv'  # Replace with the actual file path
data = irl.loadDataset(file_path)

irl.train_irl(data=data, use_dataset=True, lr=0.001, epochs=3)
irl.save_model('/Users/vinay/Desktop/Computer_Science_Projects/ReScience/hackathonf23-Stacks/models/removed_hidden_layer.pkl')