hackathonf23-Stacks/rmv_max_entropy_local.py

"""
This is the fourth ablation study, removing the maximum entropy factor during training
"""
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 IRL:
    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(2048, 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()
                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)
        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)

        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("Ablation Study 4: Removed Maximum Entropy Component")):  
            for epoch in range(epochs):
                total_loss = 0

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

                    with tf.GradientTape() as tape:
                        preferences = self.model(state[np.newaxis, :])
                        specific_loss = tf.reduce_mean(tf.square(preferences - velocity))  # Calculate Mean Squared Error

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

                avg_loss = total_loss / len(positions)
                mlflow.log_metric(f"loss", avg_loss, step=epoch)
                print(f"Epoch {epoch + 1}/{epochs}, Loss (MSE): {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:
            self.train_irl_with_dataset(data, lr=lr, epochs=epochs)
        else:
            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

            for trajectory in human_trajectories:
                for state, velocity in trajectory:  # Define velocity from the trajectory
                    with tf.GradientTape() as tape:
                        preferences = self.model(state[np.newaxis, :])
                        specific_loss = tf.reduce_mean(tf.square(preferences - velocity))  # Calculate Mean Squared Error

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

            avg_loss = total_loss / (len(human_trajectories) * trajectory_length)
            print(f"Epoch {epoch + 1}/{epochs}, Loss (MSE): {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 = IRL(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/rmv_max_entropy_model.pkl')