diff --git a/KRVI_algo_test_rotated_invariant_fulloptim.py b/KRVI_algo_test_rotated_invariant_fulloptim.py new file mode 100644 index 0000000..6ebfe0a --- /dev/null +++ b/KRVI_algo_test_rotated_invariant_fulloptim.py @@ -0,0 +1,401 @@ +from typing import Any, ClassVar, Optional, TypeVar, Union, Callable +import random +import numpy as np +import torch +import torch.nn as nn +import gymnasium as gym +from gymnasium import spaces +import wandb +import botorch +from botorch.models import SingleTaskGP +from botorch.fit import fit_gpytorch_mll +import gpytorch +from botorch.models.transforms.outcome import Standardize +from gpytorch.kernels import ScaleKernel, RBFKernel +import time +import warnings +warnings.filterwarnings("ignore") +import argparse +import os +#from gpytorch.likelihoods import GaussianLikelihood +from gpytorch.constraints import GreaterThan +import traceback + +# os.environ["WANDB_DISABLED"] = "true" +os.environ["WANDB__SERVICE_WAIT"] = "300" +device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu") +print('device',device) +#from test_randomized_env import FrozenLake2DStateWrapper +from test_rotated_reflected import FrozenLake2DStateWrapper + +from invariant_kernel_fixed import InvariantKernel, apply_rotation_group +import csv + +def action_transformation(action_index): + action_map = { + 0: np.array([-1, 0]), # Left + 1: np.array([0, -1]), # Down + 2: np.array([1, 0]), # Right + 3: np.array([0, 1]) # Up + } + return action_map.get(action_index, np.array([0, 0])) # Default to [0, 0] if invalid index + + +def preprocess_state(state): + + if isinstance(state, dict): + # Extract the observation from the dictionary + state = state['observation'] + elif isinstance(state, tuple): + # Convert the tuple to a numpy array + state = np.array(state) + elif isinstance(state, (int, float)): + # Convert scalar to a 1D numpy array + state = np.array([state]) + elif isinstance(state, np.ndarray): + # Ensure the state is a numpy array + state = state + else: + raise ValueError(f"Unsupported state type: {type(state)}") + if len(state.shape) > 1: + state = state.flatten() + + return state + + +class KRVI: + def __init__( + self, + kernel: Callable, + env: Union[gym.Env, str], + beta: float, + horizon: int, + action_transformation: Callable, + len_scale: float = 0.1, + noise_reg: float = 0.5, + optim_botorch: int = 0, + optimal_V: Optional[np.ndarray] = None, + logging: Optional[str] = None, + verbose: int = 0, + seed: Optional[int] = None, + ) -> None: + self.kernel = kernel + self.env = gym.make(env) if isinstance(env, str) else env + self.beta = beta + self.len_scale = len_scale + self.noise_reg = noise_reg + self.horizon = horizon + self.logging = logging + self.verbose = verbose + self.seed = seed + self.optim_botorch = optim_botorch + self.optimal_V = optimal_V + self.action_transformation = action_transformation + self.csv_file = 'krvi_metrics.csv' + self.config_file = 'config.txt' + + # np.random.seed(self.seed) + # torch.manual_seed(self.seed) + # torch.cuda.manual_seed(self.seed) + if self.logging: + + #wandb.init(mode='disabled') + wandb.init(project=logging, reinit=True, settings=wandb.Settings(start_method="thread")) + wandb.run.summary["noise_reg"] = self.noise_reg + wandb.run.summary["length_scale"] = self.len_scale + wandb.run.summary["UCB coef"] = self.beta + wandb.run.summary["optim_botorch"] = self.optim_botorch + wandb.run.summary["seed"] = self.seed + wandb.run.summary["kernel"]= self.kernel + + # with open(self.config_file, mode='w') as f: + # f.write(f"beta={self.beta}\n") + # f.write(f"len_scale={self.len_scale}\n") + # f.write(f"noise_reg={self.noise_reg}\n") + # f.write(f"horizon={self.horizon}\n") + # f.write(f"seed={self.seed}\n") + # f.write(f"optim_botorch={self.optim_botorch}\n") + # f.write(f"kernel={self.kernel}\n") + + # # Write headers to the metrics CSV file if it doesn't exist + # if not os.path.exists(self.csv_file): + # with open(self.csv_file, mode='w', newline='') as f: + + # writer = csv.writer(f) + # writer.writerow(['Episode', 'Reward', 'Cumulative Returns']) # Column headers + + + + def train(self, T: int): + + action_space = np.arange(self.env.action_space.n) # Assuming discrete action space + + + + # Log hyperparameters + if self.logging: + wandb.run.summary["episode length"] = self.horizon + wandb.run.summary["iterations"] = T + + # Arrays to store episode data + all_states = [] + all_actions = [] + all_rewards = [] + Qt= [None] * self.horizon + cumulative_returns = [] + + + for episode in range(T): + if self.verbose > 0: + print(f'Episode {episode}') + # Update Q-values from previous episodes + if episode > 0: + for h in reversed(range(len(all_states[-1]))): + X_states = [] + X_actions = [] + y_values = [] + + for i in range(episode): + if h < len(all_states[i]): + X_states.append(all_states[i][h]) + X_actions.append(all_actions[i][h]) + + if h < len(all_states[i]) - 1: + next_state = all_states[i][h + 1] + actions_batch = np.array([self.action_transformation(action) for action in action_space]) + states_expanded = np.tile(next_state, (len(action_space), 1)) + max_q_value = 0 + if Qt[h + 1]: + max_q_value = np.max( + self.predict_with_gp(Qt[h + 1], states_expanded, actions_batch)[0] + ) + Qnext = max_q_value + else: + Qnext = 0 + + y_values.append(all_rewards[i][h] + Qnext) + + if X_states: + X = np.column_stack((X_states, X_actions)) + y = np.array(y_values) + Qt[h] = self.GP_regression_torch(X, y) + + + # Execute episode + # Initialize arrays for the current episode + episode_states = [] + episode_actions = [] + episode_rewards = [] + + initial_state, info = self.env.reset() + + state= preprocess_state(initial_state) + + + + for h in range(self.horizon): + + if Qt[h]: # Ensure a model is available for the current step + + # Prepare inputs for batched prediction + states_batch = np.tile(state, (len(action_space), 1)) + actions_batch = np.array([self.action_transformation(action) for action in action_space]) + # Predict Q-values for all actions in a single batch + q_values = self.predict_with_gp(Qt[h], states_batch, actions_batch)[0] + + else: + # Default Q-values if no model is available + q_values = np.zeros(len(action_space)) + + # Select action with the highest Q-value + + action = action_space[np.argmax(q_values)] + next_state, reward, done, truncated , info = self.env.step(action) + next_state = preprocess_state(next_state) # Convert to numpy array + + + episode_states.append(state) + action= self.action_transformation(action) + episode_actions.append(action) + episode_rewards.append(reward) + # If done or truncated, break the loop early + if done or truncated: + break + + state = next_state + + + all_states.append(np.array(episode_states)) + all_actions.append(np.array(episode_actions)) + all_rewards.append(np.array(episode_rewards)) + + + episode_cum_rewards = np.sum(episode_rewards) + cumulative_returns.append(episode_cum_rewards) + + + if self.logging: + metrics = { + "Episode_number": episode, + "Episode_Rewards": episode_cum_rewards, + "cumulative_returns": sum(cumulative_returns) + } + wandb.log(metrics) + # with open(self.csv_file, mode='a', newline='') as f: + # writer = csv.writer(f) + # writer.writerow([episode, episode_cum_rewards, sum(cumulative_returns)]) + + + + def GP_regression_torch(self, X, y): + """ + Gaussian Process regression using PyTorch. + + :param X: Input tensor of shape (n_samples, n_features) + :param y: Target tensor of shape (n_samples,) + :return: Trained GP model + """ + X = torch.tensor(X, dtype=torch.float32,device=device) + assert not torch.isnan(X).any() + y = torch.tensor(y, dtype=torch.float32,device=device) + assert not torch.isnan(y).any() + + model = SingleTaskGP(train_X=X,train_Y= y.unsqueeze(-1).to(device)) #,outcome_transform=Standardize(m=1)) # GP expects (n_samples, 1) for targets + + model.covar_module = self.kernel.to(device) + + if isinstance(model.covar_module, gpytorch.kernels.RBFKernel): + model.covar_module.lengthscale = torch.tensor( + [self.len_scale], dtype=torch.float32, device=device + ) + + model.covar_module.raw_lengthscale.requires_grad = False + #print("The covariance module is an RBF kernel.") + else: + + # model.covar_module.base_kernel works only for the invariant kernel +# Set and freeze the length scale + model.covar_module.base_kernel.lengthscale = torch.tensor( + [self.len_scale], dtype=torch.float32, device=device + ) + + model.covar_module.base_kernel.raw_lengthscale.requires_grad = False + + # Set and freeze the noise + model.likelihood.noise = torch.tensor([self.noise_reg], dtype=torch.float32, device=device) + model.likelihood.raw_noise.requires_grad = False + if self.optim_botorch == 1: + model.likelihood.raw_noise.requires_grad = True + + try: + mll = gpytorch.mlls.ExactMarginalLogLikelihood(model.likelihood, model).to(device) + + with gpytorch.settings.cholesky_max_tries(6): + fit_gpytorch_mll(mll, optimizer_options={"n_restarts": 3, "raw_samples":60}) + del mll + + except botorch.exceptions.errors.ModelFittingError as e: + print("Model fitting failed") + + + # **Memory Cleanup** + del X, y # Safe to delete + torch.cuda.empty_cache() # Free GPU memory + + + return model + + def predict_with_gp(self, model, states_batch, actions_batch): + + X_combined = np.hstack((states_batch, actions_batch)) # Shape: (batch_size, 2) + X_combined = torch.tensor(X_combined, dtype=torch.float32, device=device) + # Make predictions + model.eval() + model.likelihood.eval() + with torch.no_grad(): + posterior = model.posterior(X_combined) + mean = posterior.mean.squeeze(-1).detach().cpu().numpy() # Shape: (batch_size,) + std_dev = posterior.variance.sqrt().squeeze(-1).detach().cpu().numpy() # Shape: (batch_size,) + # Compute mean + beta * std_dev for each batch element + acquisition_values = mean + self.beta * std_dev + return acquisition_values, mean, std_dev + + + +# Example usage +if __name__ == "__main__": + + + parser = argparse.ArgumentParser(description="Run KRVI Algorithm") + # Adding arguments for user input + parser.add_argument("--beta", type=float, default=0.1, help="UCB coefficient") + parser.add_argument("--horizon", type=int, default=100, help="Horizon length") + parser.add_argument("--len_scale", type=float, default=0.1, help="Length scale for GP kernel") + parser.add_argument("--noise_reg", type=float, default=0.1, help="Noise regularization for GP") + parser.add_argument("--env", type=str, default="FrozenLake-v1", help="Environment name") + parser.add_argument("--logging", type=str, default="trial_submission", help="wandb project name") #IQL_project_invariant + parser.add_argument("--verbose", type=int, default=1, help="Verbosity level (0: silent, 1: info)") + parser.add_argument("--iterations", type=int, default=1000, help="Number of training iterations (T)") + parser.add_argument("--seed", type=int, default=0, help="random seed") + parser.add_argument("--optim_botorch", type= int, default = 1, help ='turn on hyperparm optimization by botorch') + parser.add_argument("--kernel", type=str, default="invariant_kernel", help="Choose the kernel between invariant kernel and RBF kernel") + + + + + args = parser.parse_args() + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + random.seed(args.seed) + + env=gym.make('FrozenLake-v1', desc=None, map_name="4x4", is_slippery=False) + env = FrozenLake2DStateWrapper(env, rescale=True) + # print(env.desc) + optimal_V= None + if args.kernel=='RBF': + k_G=RBFKernel() + elif args.kernel == 'invariant_kernel': + + k_G = InvariantKernel( + base_kernel=RBFKernel(), + transformations=apply_rotation_group, + is_isotropic=True, + is_group=True, + ) + + + + krvi = KRVI( + kernel= k_G, + env= env, + beta=args.beta, + horizon=args.horizon, + action_transformation = action_transformation, + len_scale=args.len_scale, + noise_reg=args.noise_reg, + optim_botorch=args.optim_botorch, + optimal_V = optimal_V, + logging=args.logging, + verbose=args.verbose, + seed= args.seed + ) + + krvi.train(T= args.iterations) + + + + + + + + + + + + + + + + + diff --git a/KRVI_algo_test_rotated_invariant_fulloptim_eval.py b/KRVI_algo_test_rotated_invariant_fulloptim_eval.py new file mode 100644 index 0000000..a55c955 --- /dev/null +++ b/KRVI_algo_test_rotated_invariant_fulloptim_eval.py @@ -0,0 +1,442 @@ +from typing import Any, ClassVar, Optional, TypeVar, Union, Callable +import random +import numpy as np +import torch +import torch.nn as nn +import gymnasium as gym +from gymnasium import spaces +import wandb +import botorch +from botorch.models import SingleTaskGP +from botorch.fit import fit_gpytorch_mll +import gpytorch +from botorch.models.transforms.outcome import Standardize +from gpytorch.kernels import ScaleKernel, RBFKernel +import time +import warnings +warnings.filterwarnings("ignore") +import argparse +import os +#from gpytorch.likelihoods import GaussianLikelihoods +from gpytorch.constraints import GreaterThan +import traceback + +# os.environ["WANDB_DISABLED"] = "true" +os.environ["WANDB__SERVICE_WAIT"] = "300" +device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu") +print('device',device) +from test_randomized_env import FrozenLake2DStateWrapper +#from test_rotated_reflected import FrozenLake2DStateWrapper +from invariant_kernel_fixed import InvariantKernel, apply_rotation_group +import csv + +def action_transformation(action_index): + action_map = { + 0: np.array([-1, 0]), # Left + 1: np.array([0, -1]), # Down + 2: np.array([1, 0]), # Right + 3: np.array([0, 1]) # Up + } + return action_map.get(action_index, np.array([0, 0])) # Default to [0, 0] if invalid index + + +def preprocess_state(state): + + if isinstance(state, dict): + # Extract the observation from the dictionary + state = state['observation'] + elif isinstance(state, tuple): + # Convert the tuple to a numpy array + state = np.array(state) + elif isinstance(state, (int, float)): + # Convert scalar to a 1D numpy array + state = np.array([state]) + elif isinstance(state, np.ndarray): + # Ensure the state is a numpy array + state = state + else: + raise ValueError(f"Unsupported state type: {type(state)}") + if len(state.shape) > 1: + state = state.flatten() + + return state + + +class KRVI: + def __init__( + self, + kernel: Callable, + env: Union[gym.Env, str], + beta: float, + horizon: int, + action_transformation: Callable, + len_scale: float = 0.1, + noise_reg: float = 0.5, + optim_botorch: int = 0, + optimal_V: Optional[np.ndarray] = None, + logging: Optional[str] = None, + verbose: int = 0, + seed: Optional[int] = None, + ) -> None: + self.kernel = kernel + self.env = gym.make(env) if isinstance(env, str) else env + self.beta = beta + self.len_scale = len_scale + self.noise_reg = noise_reg + self.horizon = horizon + self.logging = logging + self.verbose = verbose + self.seed = seed + self.optim_botorch = optim_botorch + self.optimal_V = optimal_V + self.action_transformation = action_transformation + self.csv_file = 'krvi_metrics.csv' + self.config_file = 'config.txt' + + # np.random.seed(self.seed) + # torch.manual_seed(self.seed) + # torch.cuda.manual_seed(self.seed) + if self.logging: + + #wandb.init(mode='disabled') + wandb.init(project=logging, reinit=True, settings=wandb.Settings(start_method="thread")) + wandb.run.summary["noise_reg"] = self.noise_reg + wandb.run.summary["length_scale"] = self.len_scale + wandb.run.summary["UCB coef"] = self.beta + wandb.run.summary["optim_botorch"] = self.optim_botorch + wandb.run.summary["seed"] = self.seed + wandb.run.summary["kernel"]= self.kernel + + # with open(self.config_file, mode='w') as f: + # f.write(f"beta={self.beta}\n") + # f.write(f"len_scale={self.len_scale}\n") + # f.write(f"noise_reg={self.noise_reg}\n") + # f.write(f"horizon={self.horizon}\n") + # f.write(f"seed={self.seed}\n") + # f.write(f"optim_botorch={self.optim_botorch}\n") + # f.write(f"kernel={self.kernel}\n") + + # # Write headers to the metrics CSV file if it doesn't exist + # if not os.path.exists(self.csv_file): + # with open(self.csv_file, mode='w', newline='') as f: + + # writer = csv.writer(f) + # writer.writerow(['Episode', 'Reward', 'Cumulative Returns']) # Column headers + + + + def train(self, T: int): + + action_space = np.arange(self.env.action_space.n) # Assuming discrete action space + + + + # Log hyperparameters + if self.logging: + wandb.run.summary["episode length"] = self.horizon + wandb.run.summary["iterations"] = T + + # Arrays to store episode data + all_states = [] + all_actions = [] + all_rewards = [] + Qt= [None] * self.horizon + cumulative_returns = [] + + + for episode in range(T): + if self.verbose > 0: + print(f'Episode {episode}') + # Update Q-values from previous episodes + if episode > 0: + for h in reversed(range(len(all_states[-1]))): + X_states = [] + X_actions = [] + y_values = [] + + for i in range(episode): + if h < len(all_states[i]): + X_states.append(all_states[i][h]) + X_actions.append(all_actions[i][h]) + + if h < len(all_states[i]) - 1: + next_state = all_states[i][h + 1] + actions_batch = np.array([self.action_transformation(action) for action in action_space]) + states_expanded = np.tile(next_state, (len(action_space), 1)) + max_q_value = 0 + if Qt[h + 1]: + max_q_value = np.max( + self.predict_with_gp(Qt[h + 1], states_expanded, actions_batch)[0] + ) + Qnext = max_q_value + else: + Qnext = 0 + + y_values.append(all_rewards[i][h] + Qnext) + + if X_states: + X = np.column_stack((X_states, X_actions)) + y = np.array(y_values) + Qt[h] = self.GP_regression_torch(X, y) + + + # Execute episode + # Initialize arrays for the current episode + episode_states = [] + episode_actions = [] + episode_rewards = [] + + initial_state, info = self.env.reset() + + + state= preprocess_state(initial_state) + + + + for h in range(self.horizon): + + if Qt[h]: # Ensure a model is available for the current step + + # Prepare inputs for batched prediction + states_batch = np.tile(state, (len(action_space), 1)) + actions_batch = np.array([self.action_transformation(action) for action in action_space]) + # Predict Q-values for all actions in a single batch + q_values = self.predict_with_gp(Qt[h], states_batch, actions_batch)[0] + + else: + # Default Q-values if no model is available + q_values = np.zeros(len(action_space)) + + # Select action with the highest Q-value + + action = action_space[np.argmax(q_values)] + next_state, reward, done, truncated , info = self.env.step(action) + next_state = preprocess_state(next_state) # Convert to numpy array + + + episode_states.append(state) + action= self.action_transformation(action) + episode_actions.append(action) + episode_rewards.append(reward) + # If done or truncated, break the loop early + if done or truncated: + break + + state = next_state + + + all_states.append(np.array(episode_states)) + all_actions.append(np.array(episode_actions)) + all_rewards.append(np.array(episode_rewards)) + + + episode_cum_rewards = np.sum(episode_rewards) + cumulative_returns.append(episode_cum_rewards) + eval_score, eval_std = self.evaluate(Qt, action_space, num_episodes=40) + print(f"Evaluation after episode {episode}: Avg return = {eval_score:.2f}") + + if self.logging: + metrics = { + "Episode_number": episode, + "Episode_Rewards": episode_cum_rewards, + "cumulative_returns": sum(cumulative_returns), + "Evaluation_Avg_Return": eval_score, + "Evaluation_Std_Dev": eval_std + + } + wandb.log(metrics) + # with open(self.csv_file, mode='a', newline='') as f: + # writer = csv.writer(f) + # writer.writerow([episode, episode_cum_rewards, sum(cumulative_returns)]) + + + + + def GP_regression_torch(self, X, y): #I removed normalization + """ + Gaussian Process regression using PyTorch. + + :param X: Input tensor of shape (n_samples, n_features) + :param y: Target tensor of shape (n_samples,) + :return: Trained GP model + """ + X = torch.tensor(X, dtype=torch.float32,device=device) + assert not torch.isnan(X).any() + y = torch.tensor(y, dtype=torch.float32,device=device) + assert not torch.isnan(y).any() + + model = SingleTaskGP(train_X=X,train_Y= y.unsqueeze(-1).to(device)) #,outcome_transform=Standardize(m=1)) # GP expects (n_samples, 1) for targets + + model.covar_module = self.kernel.to(device) + + if isinstance(model.covar_module, gpytorch.kernels.RBFKernel): + model.covar_module.lengthscale = torch.tensor( + [self.len_scale], dtype=torch.float32, device=device + ) + + model.covar_module.raw_lengthscale.requires_grad = False + #print("The covariance module is an RBF kernel.") + else: + + # model.covar_module.base_kernel works only for the invariant kernel +# Set and freeze the length scale + model.covar_module.base_kernel.lengthscale = torch.tensor( + [self.len_scale], dtype=torch.float32, device=device + ) + + model.covar_module.base_kernel.raw_lengthscale.requires_grad = False + + # Set and freeze the noise + model.likelihood.noise = torch.tensor([self.noise_reg], dtype=torch.float32, device=device) + model.likelihood.raw_noise.requires_grad = False + if self.optim_botorch == 1: + model.likelihood.raw_noise.requires_grad = True + + try: + mll = gpytorch.mlls.ExactMarginalLogLikelihood(model.likelihood, model).to(device) + + with gpytorch.settings.cholesky_max_tries(6): + fit_gpytorch_mll(mll, optimizer_options={"n_restarts": 3, "raw_samples":60}) + del mll + + except botorch.exceptions.errors.ModelFittingError as e: + print("Model fitting failed") + + + # **Memory Cleanup** + del X, y # Safe to delete + torch.cuda.empty_cache() # Free GPU memory + + + return model + + def predict_with_gp(self, model, states_batch, actions_batch): + + X_combined = np.hstack((states_batch, actions_batch)) # Shape: (batch_size, 2) + X_combined = torch.tensor(X_combined, dtype=torch.float32, device=device) + # Make predictions + model.eval() + model.likelihood.eval() + with torch.no_grad(): + posterior = model.posterior(X_combined) + mean = posterior.mean.squeeze(-1).detach().cpu().numpy() # Shape: (batch_size,) + std_dev = posterior.variance.sqrt().squeeze(-1).detach().cpu().numpy() # Shape: (batch_size,) + # Compute mean + beta * std_dev for each batch element + acquisition_values = mean + self.beta * std_dev + return acquisition_values, mean, std_dev + + def evaluate(self, Qt, action_space, num_episodes=40): #evaluating on 40 different environments + test_returns = [] + + test_env = FrozenLake2DStateWrapper(gym.make("FrozenLake-v1", is_slippery=False), rescale=True) + + for _ in range(num_episodes): + state, _ = test_env.reset() + state = preprocess_state(state) + + episode_reward = [] + total_reward = 0 + + for h in range(self.horizon): + if Qt[h]: + states_batch = np.tile(state, (len(action_space), 1)) + actions_batch = np.array([self.action_transformation(a) for a in action_space]) + q_values = self.predict_with_gp(Qt[h], states_batch, actions_batch)[0] + else: + q_values = np.zeros(len(action_space)) + + action = action_space[np.argmax(q_values)] + next_state, reward, done, truncated, _ = test_env.step(action) + next_state = preprocess_state(next_state) + + total_reward += reward + + if done or truncated: + break + + state = next_state + + test_returns.append(total_reward) + + return np.mean(test_returns), np.std(test_returns) + + + + +# Example usage +if __name__ == "__main__": + + + parser = argparse.ArgumentParser(description="Run KRVI Algorithm") + # Adding arguments for user input + parser.add_argument("--beta", type=float, default=0.1, help="UCB coefficient") + parser.add_argument("--horizon", type=int, default=100, help="Horizon length") + parser.add_argument("--len_scale", type=float, default=0.1, help="Length scale for GP kernel") + parser.add_argument("--noise_reg", type=float, default=0.1, help="Noise regularization for GP") + parser.add_argument("--env", type=str, default="FrozenLake-v1", help="Environment name") + parser.add_argument("--logging", type=str, default="trial_submission", help="wandb project name") + parser.add_argument("--verbose", type=int, default=1, help="Verbosity level (0: silent, 1: info)") + parser.add_argument("--iterations", type=int, default=2000, help="Number of training iterations (T)") + parser.add_argument("--seed", type=int, default=0, help="random seed") + parser.add_argument("--optim_botorch", type= int, default = 1, help ='turn on hyperparm optimization by botorch') + parser.add_argument("--kernel", type=str, default="invariant_kernel", help="Choose the kernel between invariant kernel and RBF kernel") + + + + + args = parser.parse_args() + np.random.seed(args.seed) + torch.manual_seed(args.seed) + torch.cuda.manual_seed(args.seed) + random.seed(args.seed) + + env=gym.make('FrozenLake-v1', desc=None, map_name="4x4", is_slippery=False) + env = FrozenLake2DStateWrapper(env, rescale=True) + # print(env.desc) + optimal_V= None + if args.kernel=='RBF': + k_G=RBFKernel() + elif args.kernel == 'invariant_kernel': + + k_G = InvariantKernel( + base_kernel=RBFKernel(), + transformations=apply_rotation_group, + is_isotropic=True, + is_group=True, + ) + + + + krvi = KRVI( + kernel= k_G, + env= env, + beta=args.beta, + horizon=args.horizon, + action_transformation = action_transformation, + len_scale=args.len_scale, + noise_reg=args.noise_reg, + optim_botorch=args.optim_botorch, + optimal_V = optimal_V, + logging=args.logging, + verbose=args.verbose, + seed= args.seed + ) + + krvi.train(T= args.iterations) + + + + + + + + + + + + + + + + + diff --git a/README.md b/README.md index 72161e4..6bfea7a 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,47 @@ # IQL -Kernelized Q-learning with invariances. +Kernelized Q-learning with invariances (IQL) is an implementation of LSVI algorithm designed to incorporate invariance properties into the kernel function. This can improve sample efficiency and generalization, especially in structured environments. + +## Synthetic Experiments +To run synthetic experiments: + +```bash +python script.py +``` +It will run the KRVI algorithm with the invariant kernel on the synthetically generated reward and transition dynamics functions, for T=1000 episodes, using the UCB coefficient $\beta$= 0.1 and averaging over NUM_RUNS=20 runs. You can change these parameters in script.py +## Frozen Lake - Fixed Layout +To run experiments on the Frozen Lake environment with a fixed layout: + +```bash +python KRVI_algo_test_rotated_invariant_fulloptim.py +``` +You can specify hyperparparameters as arguments. + +We used the following hyperparameters in our experiments: + +For the invariant kernel: +```bash + --kernel invariant_kernel --beta 0.01 --len_scale 0.5 --noise_reg 0.1 --optim_botorch 1 --iterations 1500 + ``` + +For the RBF kernel: +```bash + --kernel RBF --beta 0.01 --len_scale 0.1 --noise_reg 0.1 --optim_botorch 1 --iterations 1500 + ``` +## Frozen Lake - Random Layout +To run experiments on the Frozen Lake environment with random layouts: + + +```bash +python KRVI_algo_test_rotated_invariant_fulloptim_eval.py +``` +We used the following hyperparameters in our experiments: + +For the invariant kernel: +```bash +--kernel invariant_kernel --beta 0.01 --len_scale 0.5 --noise_reg 0.1 --optim_botorch 1 --iterations 2000 +``` + +For the RBF kernel: +```bash + --kernel RBF --beta 0.01 --len_scale 1 --noise_reg 0.1 --optim_botorch 1 --iterations 2000 +``` \ No newline at end of file diff --git a/common_framework2.py b/common_framework2.py new file mode 100644 index 0000000..a28a604 --- /dev/null +++ b/common_framework2.py @@ -0,0 +1,36 @@ +from framework import transition_P_RKHS, reward_RKHS, value_iteration_episodic, plot_reward_gp_3d, plot_transition_probabilities +import numpy as np +import os +from datetime import datetime +import matplotlib.pyplot as plt + +def save_rewards_and_transitions(P_kernel, alpha, save_dir): + state_space = np.linspace(-1, 1, num=10).reshape(-1, 1) + action_space = np.linspace(-1, 1, num=10).reshape(-1, 1) + H = 10 + + subdir = os.path.join(save_dir, P_kernel) + os.makedirs(subdir, exist_ok=True) + + P = transition_P_RKHS(state_space, action_space, P_kernel, alpha) + r = reward_RKHS(P_kernel, state_space, action_space, subdir, alpha) + optimal_value_function = value_iteration_episodic(state_space, action_space, r, P, H) + + np.save(os.path.join(subdir, 'P.npy'), P) + np.save(os.path.join(subdir, 'r.npy'), r) + np.save(os.path.join(subdir, 'V.npy'), optimal_value_function) + plot_reward_gp_3d(r, state_space, action_space, save_dir) + plot_transition_probabilities(P, state_space,action_space, save_dir) + plt.plot(optimal_value_function) + plt.ylabel("optimal value function") + save_path = os.path.join(save_dir, f"optimal_V.png") + plt.savefig(save_path) + + + +def load_saved_data(subdir): + P = np.load(os.path.join(subdir, 'P.npy')) + r = np.load(os.path.join(subdir, 'r.npy')) + optimal_value_function = np.load(os.path.join(subdir, 'V.npy')) + return P, r, optimal_value_function + diff --git a/framework.py b/framework.py new file mode 100644 index 0000000..cf61557 --- /dev/null +++ b/framework.py @@ -0,0 +1,287 @@ +import numpy as np +import matplotlib.pyplot as plt +from sklearn.kernel_ridge import KernelRidge +from sklearn.gaussian_process import GaussianProcessRegressor +from sklearn.gaussian_process.kernels import RBF, Matern, Kernel +from scipy.stats import truncnorm +import os +import wandb +import datetime +from itertools import product +from mpl_toolkits.mplot3d import Axes3D +from sklearn.metrics import mean_squared_error +import datetime +from itertools import product + +# #kernel implementation from Scikit learn +class GroupInvariantKernel(Kernel): + def __init__(self, base_kernel="Matern", length_scale=0.1, smoothness=1.5, group=None): + self.base_kernel = base_kernel + self.length_scale = length_scale + self.smoothness = smoothness + self.group = group if group is not None else [lambda x: x] # Default to identity transformation if no group given + + # Define the base kernel as Matern or RBF + if self.base_kernel == "Matern": + self.kernel = Matern(length_scale=length_scale, nu=smoothness) + elif self.base_kernel == "RBF": + self.kernel = RBF(length_scale=length_scale) + + def __call__(self, X, Y=None, eval_gradient=False): + return self.group_invariant_kernel(X, Y,eval_gradient=eval_gradient) + + def group_invariant_kernel(self, X, Y=None, eval_gradient=False): + if Y is None: + Y = X + X = np.atleast_2d(X) + # print('X',X) + Y = np.atleast_2d(Y) + # Ensure the dimensions are as expected + + # print('Y',Y) + # Calculate the kernel averaged over the group transformations + K = np.zeros((X.shape[0], Y.shape[0])) + + # # Loop over all transformations in the group + for g in self.group: + #print('g',g) + X_transformed = np.array([g(x) for x in X]) + #print('X_transformed',X_transformed.shape) + # print('X_transformed values',X_transformed) + # print('Y',Y.shape) + # Check if the kernel is producing a valid square matrix for each transformation + transformed_kernel = self.kernel(X_transformed, Y) + #print('Transformed kernel shape:', transformed_kernel.shape) + + K += transformed_kernel / len(self.group) + + + + return K + + def diag(self, X): + return np.diag(self.__call__(X)) + + + def is_stationary(self): + return True + #return self.kernel.is_stationary() +# +def shift_05(x): + return (x + 0.5) % 1 + +def shift_03(x): + return (x + 0.3) % 1 +def square(x): + return x**2 + +def flip(x): + return - x + + +group_SA = [ + lambda x: np.array([x[0], x[1]]), # Identity for both state and action + lambda x: np.array([flip(x[0]), flip(x[1])]) # Flip for both state and action +] + + + + +group_SASp = [ + lambda x: np.array([x[0],x[1],x[2]]), # Identity for both state and action + lambda x: np.array([flip(x[0]), flip(x[1]), flip(x[2])]) # Flip for both state and action +] + + + +# Define the reward_RKHS function with the group-invariant kernel +def reward_RKHS(P_kernel, state_space, action_space, subdir=None, alpha=0.5): + grid_size = 5 # Grid size for fitting GP regression + + # Generate all possible input points in the grid + values = np.linspace(-1, 1, grid_size) + X = np.array(list(product(values, repeat=2))) # 2D grid points for state-action pairs + + # Define the group-invariant kernel based on the P_kernel parameter and group G (state-action transformations) + if P_kernel == "Matern_smoothness_1.5": + kernel = GroupInvariantKernel(base_kernel="Matern", length_scale=0.1, smoothness=1.5, group=group_SA) + elif P_kernel == "Matern_smoothness_2.5": + kernel = GroupInvariantKernel(base_kernel="Matern", length_scale=0.1, smoothness=2.5, group=group_SA) + + elif P_kernel == "RBF": + kernel = GroupInvariantKernel(base_kernel="RBF", length_scale=0.1, group=group_SA) + + # Sample y values from the Gaussian process with the group-invariant kernel + gp = GaussianProcessRegressor(kernel=kernel) + + y = gp.sample_y(X, 1) + + gpr = GaussianProcessRegressor(kernel=kernel, optimizer=None, alpha=alpha) + K = gpr.kernel(X,X) + eigvals = np.linalg.eigvalsh(K) + + gpr.fit(X, y) + + values = np.linspace(-1, 1, len(state_space)) + all_possible_inputs = np.array(list(product(values, repeat=2))) # State-action pairs in 2D + + # Predict for all possible input points + all_predictions, _ = gpr.predict(all_possible_inputs, return_std=True) + y_pred, _ = gpr.predict(X, return_std=True) + mse = mean_squared_error(y, y_pred) + # Scale and normalize the predictions + min_prediction = np.min(all_predictions) + max_prediction = np.max(all_predictions) + scaled_predictions = (all_predictions - min_prediction) / (max_prediction - min_prediction) + r = scaled_predictions.reshape((len(state_space), len(action_space))) + + return r + +def transition_P_RKHS(state_space, action_space, P_kernel, alpha=0.5): + grid_size = 5 # Grid size for fitting GP regression + values = np.linspace(-1, 1, grid_size) + X = np.array(list(product(values, repeat=3))) # 3D points for state-action-next_state + + # # Define the group-invariant kernel based on the P_kernel parameter and transformations for 3D inputs + if P_kernel == "Matern_smoothness_1.5": + kernel = GroupInvariantKernel(base_kernel="Matern", length_scale=0.1, smoothness=1.5, group=group_SASp) + elif P_kernel == "Matern_smoothness_2.5": + kernel = GroupInvariantKernel(base_kernel="Matern", length_scale=0.1, smoothness=2.5, group=group_SASp) + elif P_kernel == "RBF": + kernel = GroupInvariantKernel(base_kernel="RBF", length_scale=0.1, group=group_SASp) + + # Sample y values from the Gaussian process + gp = GaussianProcessRegressor(kernel=kernel) + y = gp.sample_y(X, 1) + + # Fit the Gaussian Process Regressor + gpr = GaussianProcessRegressor(kernel=kernel, optimizer=None, alpha=alpha) + gpr.fit(X, y) + print('after fitting') + + # Generate all possible input points for prediction (state-action-next_state) + values = np.linspace(-1, 1, len(state_space)) + all_possible_inputs = np.array(list(product(values, repeat=3))) # 3D points for state-action-next_state + + # Predict for all possible input points + all_predictions, _ = gpr.predict(all_possible_inputs, return_std=True) + + # Scale up and normalize the predictions + min_prediction = np.min(all_predictions) + max_prediction = np.max(all_predictions) + scaled_predictions = (all_predictions - min_prediction) / (max_prediction - min_prediction) + + num_state_action_pairs = len(all_possible_inputs) // len(state_space) + for i in range(num_state_action_pairs): + state_start_index = i * len(state_space) + state_end_index = (i + 1) * len(state_space) + state_predictions = scaled_predictions[state_start_index:state_end_index] + scaled_predictions[state_start_index:state_end_index] = state_predictions / np.sum(state_predictions) + + # Reshape the predictions to form the transition probability tensor P + P = scaled_predictions.reshape((len(state_space), len(action_space), len(state_space))) + print('P,', P) + return P + + +# Define transition dynamics function which takes current state and action values as input and outputs the value of the next state +def transition_dynamics(current_state, action, transition_P,state_space,action_space): + + # Get the index of the current state + current_state_idx = np.argmin(np.abs(state_space - current_state)) + # Get the index of the action + action_idx = np.argmin(np.abs(action_space - action)) + + # Retrieve transition probabilities from the precomputed dynamics + transition_probs = transition_P[current_state_idx, action_idx, :] + + # Sample next state based on transition probabilities + next_state_idx = np.random.choice(np.arange(len(state_space)), p=transition_probs) + + next_state = state_space[next_state_idx] + + return next_state + + +# Value iteration algorithm +def value_iteration_episodic(state_space, action_space, r, P, H=10): + V = np.zeros_like(state_space) # Initialize value function + for h in range(H): # Iterate over time steps + V_new = np.zeros_like(V) # Initialize new value function + for s_idx, s in enumerate(state_space): + Q_values = [] + for a_idx, a in enumerate(action_space): + expected_return = 0 + for next_s_idx, next_s in enumerate(state_space): + reward = r[s_idx,a_idx] + expected_return += P[s_idx, a_idx, next_s_idx] * (reward + V[next_s_idx]) + Q_values.append(expected_return) + V_new[s_idx] = max(Q_values) # Update value function for state s + V = V_new # Update value function after each iteration + + return V # V is indexed by state index + + +def plot_reward_gp_3d(r, state_space, action_space, save_dir=None): + fig = plt.figure(figsize=(10, 8)) + ax = fig.add_subplot(111, projection='3d') + # Plot the surface + action_mesh, state_mesh = np.meshgrid(state_space, action_space) + surf = ax.plot_surface(state_mesh, action_mesh, r, cmap='viridis',vmin=0,vmax=1) + fig.colorbar(surf, shrink=0.5, aspect=5) + + + + ax.set_xlabel('s',fontsize=20) + ax.set_ylabel('a',fontsize=20) + ax.set_zlabel('r(s,a)',fontsize=20) + + if save_dir: + if not os.path.exists(save_dir): + os.makedirs(save_dir) + timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + save_path = os.path.join(save_dir, f"reward_gp_3d_{timestamp}.png") + plt.savefig(save_path) + plt.close() + else: + wandb.log({"Reward Surface 3D": wandb.Image(plt)}) + + + +def plot_transition_probabilities(P, state_space,action_space, save_dir=None): + + + + state_indices = [0, 5, 9] # Indices of states in state_space + action_indices = [0, 5, 9] # Indices of actions in action_space + +# Plot transition probabilities for each (s, a) pair + for state_idx in state_indices: + for action_idx in action_indices: + # Get the transition probabilities for the specified state-action pair + transition_probs = P[state_idx, action_idx, :] + + state_labels = [str(round(float(state[0]), 4)) for state in state_space] + action_labels = [str(round(float(action[0]), 4)) for action in action_space] + + # Plot + plt.figure(figsize=(8, 6)) + plt.bar(state_labels, transition_probs) + plt.xlabel("s'",fontsize=20) + # plt.rcParams['text.usetex'] = True + # plt.rcParams['text.latex.preamble'] = r'\usepackage{{amsmath}}' + + plt.ylabel(r"$P(s' \mid s={}, a={})$".format(state_labels[state_idx], action_labels[action_idx]), fontsize=20) + + plt.xticks([state_labels[0], state_labels[-1]]) + + if save_dir: + if not os.path.exists(save_dir): + os.makedirs(save_dir) + save_path = os.path.join(save_dir, f"P_s{state_idx}_a{action_idx}.png") + plt.savefig(save_path) # Save the figure with a unique name + plt.close() + else: + # Convert the Matplotlib figure to a wandb Image and log it + wandb.log({f"Transition Probabilities for State-Action Pair (s={state_space[state_idx]}, a={action_space[action_idx]})": wandb.Image(plt)}) + diff --git a/invariant_kernel_fixed.py b/invariant_kernel_fixed.py new file mode 100644 index 0000000..746a536 --- /dev/null +++ b/invariant_kernel_fixed.py @@ -0,0 +1,222 @@ +import itertools +from typing import Iterator +import torch +import gpytorch +from gpytorch.kernels import MaternKernel, ScaleKernel +from typing import Callable, Iterable, Optional, Tuple +import torch +import numpy as np +from itertools import product +from gpytorch.means import ZeroMean +from gpytorch.kernels import MaternKernel, RBFKernel +from gpytorch.likelihoods import GaussianLikelihood +from gpytorch.mlls import ExactMarginalLogLikelihood +import torch +import numpy as np +from scipy.linalg import block_diag + + + +class InvariantKernel(gpytorch.kernels.Kernel): + r"""A kernel that is invariant to a collection of transformations. + + Currently only supports group invariance. + + The invariant kernel is defined as: + .. math:: + k_G(x, y) = \frac{1}{|G|^2} \sum_{g \in G} \sum_{h \in G} k(g(x), h(y)) + + If the kernel is isotropic, we can use the simpler form: + .. math:: + k_G(x, y) = \frac{1}{|G|} \sum_{g \in G} k(g(x), y) + + where :math:`G` is the group of transformations, :math:`k` is the base kernel, and :math:`x` and :math:`y` are the inputs. + """ + + def __init__( + self, + base_kernel: gpytorch.kernels.Kernel, + transformations: Callable[[torch.tensor], torch.tensor], + is_isotropic: bool = True, + is_group: bool = True, + **kwargs, + ) -> None: + super().__init__(**kwargs) + + self.base_kernel = base_kernel + self.transformations = transformations + self.is_isotropic = is_isotropic + self.is_group = is_group + + if not self.is_group: + raise NotImplementedError("InvariantKernel only supports group invariance.") + + def forward( + self, x1, x2, diag: bool = False, last_dim_is_batch: bool = False, **kwargs + ) -> torch.tensor: + if last_dim_is_batch: + raise NotImplementedError( + "last_dim_is_batch=True not supported for GroupInvariantKernel." + ) + # print('x1',x1) + + x1_orbits = self.transformations(x1) # Shape is ... x G x N x d + # print('x1_orbits',x1_orbits) + #print('x1_orbits shape',x1_orbits.shape) + G = x1_orbits.shape[-3] + if self.is_isotropic: + # Sum is over a single set of orbits + # x2_orbits should be constructed by tiling x2 along the -3 axis + dims = [-1] * (x2.dim() + 1) + #print('dims',dims) + dims[-3] = G + #print('x2',x2) + x2_orbits = x2.unsqueeze(-3).expand(dims) + #print('x2_orbits',x2_orbits) + #print('x2_orbits',x2_orbits.shape) + K_orbits = self.base_kernel.forward(x1_orbits, x2_orbits) + #print('K_orbits',K_orbits.shape) + K = torch.mean(K_orbits, dim=-3) + # print("Grad check x1_orbits requires_grad:", x1_orbits.requires_grad) + # print("Grad check x2_orbits requires_grad:", x2_orbits.requires_grad) + + #print('K',K.shape) + else: + # Sum is over all pairs of orbits + # x2_orbits should be constructed by applying the transformations to x2 + x2_orbits = self.transformations(x2) # Shape is ... x G x M x d + + if x2_orbits.shape[-3] != G or x1_orbits.shape[-3] != G: + raise ValueError( + "Different numbers of orbits for x1 and x2. " + "Check that self.transformations returns a tensor of shape (..., G, N, d)." + ) + + # WARNING: This is quadratic in the number of orbits! + # TODO: We should be able to parallelise this + # TODO: Should be a way of doing it with less memory + # Repeat each element of x1_orbits G times + # New shape is G^2 x ... x N x d + x1_orbits_expanded = x1_orbits.repeat_interleave(G, dim=-3) + # Repeat the entire x2_orbits G times + # New shape is ... x G^2 x M x d + dims = [1] * x2_orbits.dim() + dims[-3] = G + x2_orbits_expanded = x2_orbits.repeat(dims) + # Compute the kernel between each pair of expanded orbits = all combinations of orbits + K_orbits = self.base_kernel.forward(x1_orbits_expanded, x2_orbits_expanded) + K = torch.mean(K_orbits, dim=-3) + #print('kernel',K) + + if diag: + return K.diag() + else: + return K + + + + + +def construct_rot_and_reflection_group(dim_space: int): + """ + Constructs a group of 8 transformations: + - 4 rotations (0°, 90°, 180°, 270°) + - 4 rotations followed by reflection across the x-axis + """ + assert dim_space % 2 == 0, "Dimension must be even (pairs of x, y)." + + # 2D 90-degree rotation matrix (counter-clockwise) + rot_90 = np.array([[0, 1], [-1, 0]]) + rot_blocks = [rot_90] * (dim_space // 2) + rot_matrix = block_diag(*rot_blocks) + + # 2D reflection over x-axis + reflect_x = np.array([[1, 0], [0, -1]]) + reflect_blocks = [reflect_x] * (dim_space // 2) + reflect_matrix_x = block_diag(*reflect_blocks) + + # Generate rotation matrices + rotations = [np.eye(dim_space)] + current = np.eye(dim_space) + for _ in range(3): + current = current @ rot_matrix + rotations.append(current.copy()) + + # Now: rotations followed by reflection in x-axis + rot_reflections = [R @ reflect_matrix_x for R in rotations] + + # Combine all 8 + full_group = rotations + rot_reflections + + return [torch.tensor(M, dtype=torch.float32) for M in full_group] + + + + + +def apply_rotation_group(x: torch.Tensor) -> torch.Tensor: + """ + Applies 8 symmetry transformations: 4 rotations + 4 rot-reflections. + + Args: + x: Tensor of shape (..., d), where d = 2 * num_pairs. + + Returns: + Tensor of shape (..., 8, d), one per transformation. + """ + device = x.device + group = construct_rot_and_reflection_group(x.shape[-1]) + group = [g.to(device=device, dtype=torch.float32) for g in group] + transformed = [x @ g.T for g in group] + return torch.stack(transformed, dim=-3) + +def is_positive_definite(K: torch.Tensor, tol: float = 1e-5) -> bool: + """ + Check if a kernel matrix is positive semi-definite by checking if all its eigenvalues + are non-negative. + + Args: + K (torch.Tensor): The kernel matrix to check. + tol (float): The tolerance to consider eigenvalues as non-negative. + + Returns: + bool: True if the matrix is positive semi-definite, otherwise False. + """ + # Eigenvalues of the kernel matrix + eigenvalues = torch.linalg.eigvals(K) + + # Check if all eigenvalues are non-negative + return torch.all(eigenvalues.real >= -tol) + + +def main(): + dim_space = 14 # 14-dimensional vector + # print( construct_rot_and_reflection_group(dim_space=14)) + n_datapoints = 10 + + torch.manual_seed(0) + x = torch.rand((n_datapoints, dim_space)) + + k = ScaleKernel(MaternKernel(nu=2.5)) + + k_G = InvariantKernel( + base_kernel=k, + transformations=apply_rotation_group, + is_isotropic=True, + is_group=True, + ) + + # manual_k_G_matrix = compute_manual_isotropic_kernel_matrix(k, x, x, apply_rotation_group) + # print("Manual kernel matrix:\n", manual_k_G_matrix) + + k_G_matrix = k_G(x, x).to_dense() + print("Invariant kernel matrix:\n", k_G_matrix) + + # assert torch.allclose(k_G_matrix, manual_k_G_matrix, atol=1e-5) + if is_positive_definite( k_G_matrix): + print("The kernel matrix is positive definite.") + else: + print("The kernel matrix is NOT positive definite.") +if __name__ == "__main__": + main() + diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..edb07ac --- /dev/null +++ b/requirements.txt @@ -0,0 +1,77 @@ +botorch==0.10.0 +certifi==2024.8.30 +charset-normalizer==3.4.0 +click==8.1.7 +cloudpickle==3.1.0 +contourpy==1.3.0 +cycler==0.12.1 +docker-pycreds==0.4.0 +Farama-Notifications==0.0.4 +filelock==3.16.1 +fonttools==4.54.1 +fsspec==2024.10.0 +gitdb==4.0.11 +GitPython==3.1.43 +gpytorch==1.11 +gym==0.26.2 +gym-notices==0.0.8 +gymnasium==1.0.0 +idna==3.10 +importlib_metadata==8.5.0 +importlib_resources==6.4.5 +jaxtyping==0.2.19 +Jinja2==3.1.4 +joblib==1.4.2 +kiwisolver==1.4.7 +linear-operator==0.5.1 +MarkupSafe==3.0.2 +matplotlib==3.9.2 +minigrid==3.0.0 +mpmath==1.3.0 +multipledispatch==1.0.0 +networkx==3.2.1 +numpy==2.0.2 +nvidia-cublas-cu12==12.4.5.8 +nvidia-cuda-cupti-cu12==12.4.127 +nvidia-cuda-nvrtc-cu12==12.4.127 +nvidia-cuda-runtime-cu12==12.4.127 +nvidia-cudnn-cu12==9.1.0.70 +nvidia-cufft-cu12==11.2.1.3 +nvidia-curand-cu12==10.3.5.147 +nvidia-cusolver-cu12==11.6.1.9 +nvidia-cusparse-cu12==12.3.1.170 +nvidia-nccl-cu12==2.21.5 +nvidia-nvjitlink-cu12==12.4.127 +nvidia-nvtx-cu12==12.4.127 +opt_einsum==3.4.0 +packaging==24.2 +pandas==2.2.3 +pillow==11.0.0 +platformdirs==4.3.6 +protobuf==5.28.3 +psutil==6.1.0 +pygame==2.6.1 +pyparsing==3.2.0 +pyro-api==0.1.2 +pyro-ppl==1.9.1 +python-dateutil==2.9.0.post0 +pytz==2024.2 +PyYAML==6.0.2 +requests==2.32.3 +scikit-learn==1.5.2 +scipy==1.13.1 +sentry-sdk==2.18.0 +setproctitle==1.3.3 +six==1.16.0 +smmap==5.0.1 +sympy==1.13.1 +threadpoolctl==3.5.0 +torch==2.5.1 +tqdm==4.67.1 +triton==3.1.0 +typeguard==2.13.3 +typing_extensions==4.12.2 +tzdata==2024.2 +urllib3==2.2.3 +wandb==0.18.6 +zipp==3.21.0 diff --git a/script.py b/script.py new file mode 100644 index 0000000..df78fdf --- /dev/null +++ b/script.py @@ -0,0 +1,57 @@ +import test_kernelized_RL_invariant +from common_framework2 import save_rewards_and_transitions, load_saved_data +import numpy as np +import wandb +import os +from datetime import datetime + +def main(): + #Define state and action spaces and their combination + state_space = np.linspace(-1, 1, num=10).reshape(-1, 1) + action_space = np.linspace(-1, 1, num=10).reshape(-1, 1) + + # Define state-action space + state_action_space = np.array([np.hstack((state, action)) for state in state_space for action in action_space]) + # Define environment parameters + S = len(state_space) + A = len(action_space) + H = 10 # Length of each episode + + P_kernel= 'RBF' + alpha=0.01 #0.001 + # Define save directory + #save_dir = f'Rewards_and_P_invariant_NEW' + #subdir = os.path.join(save_dir, P_kernel) + # Add a timestamp to the save directory + timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + save_dir = f'Rewards_and_P_invariant_NEW_{alpha}' + + subdir = os.path.join(save_dir, P_kernel) + + if not os.path.exists(subdir) or not os.listdir(subdir): + print('not there') + save_rewards_and_transitions(P_kernel, alpha, save_dir) + + P, r, optimal_value_function = load_saved_data(subdir) + + + + # Define environment M = (S,A,H,P,r) + M = (S, A, H, P, r) + + # Run experiments with p-KRVI policy + T = 1000 # Number of episodes + beta=0.1 # UCB coefficient (0.001) + NUM_RUNS=20 + + for run in range(NUM_RUNS): + # Start a new run for each iteration + wandb.init(project="IQL_submission_test", reinit=True) + wandb.run.summary["alpha_functions"] = alpha + test_kernelized_RL_invariant.pi_krvi_policy(M, T, state_space, action_space,state_action_space,optimal_value_function, beta) + # Log metrics specific to this run + wandb.log({"run": run}) + +if __name__ == "__main__": + main() + diff --git a/test_kernelized_RL_invariant.py b/test_kernelized_RL_invariant.py new file mode 100644 index 0000000..7992730 --- /dev/null +++ b/test_kernelized_RL_invariant.py @@ -0,0 +1,166 @@ +import numpy as np +import matplotlib.pyplot as plt +from sklearn.kernel_ridge import KernelRidge +from sklearn.gaussian_process import GaussianProcessRegressor +from sklearn.gaussian_process.kernels import RBF, Matern +from scipy.stats import truncnorm +import wandb +import datetime +from framework import transition_dynamics, GroupInvariantKernel + +def flip(x): + return - x +group_SA = [ + lambda x: np.array([x[0], x[1]]), # Identity for both state and action + lambda x: np.array([flip(x[0]), flip(x[1])]) # Flip for both state and action +] + +# Function for GP regression to estimate Q-values using RBF kernel +def GP_regression_invariant(X, y, state_action_space,P_kernel='RBF',l=1,alpha=1e-10): + + if P_kernel == "Matern_smoothness_1.5": + kernel = GroupInvariantKernel(base_kernel="Matern", length_scale=l, smoothness=1.5, group=group_SA) + elif P_kernel == "Matern_smoothness_2.5": + kernel = GroupInvariantKernel(base_kernel="Matern", length_scale=l, smoothness=2.5, group=group_SA) + + elif P_kernel == "RBF": + kernel = GroupInvariantKernel(base_kernel="RBF", length_scale=l, group=group_SA) + + # Define the RBF kernel + + wandb.run.summary["kernel type"] = kernel + wandb.run.summary["alpha_gp"] = alpha + wandb.run.summary["length_scale"] = l + # Create GPR model + gpr = GaussianProcessRegressor(kernel=kernel, optimizer=None,alpha=alpha) + # Fit the model + gpr.fit(X, y) + # Predict mean and standard deviation + y_pred_mean, y_pred_std = gpr.predict(state_action_space, return_std=True) + + return y_pred_mean, y_pred_std + + +# Function for GP regression to estimate Q-values using RBF kernel +def GP_regression_with_RBF(X, y, state_action_space,l=1,alpha=1e-10): + # Define the RBF kernel + kernel = RBF(length_scale=l,length_scale_bounds="fixed") + + #tau=0.01 + wandb.run.summary["kernel type"] = kernel + wandb.run.summary["alpha_gp"] = alpha + wandb.run.summary["length_scale"] = l + # Create GPR model + gpr = GaussianProcessRegressor(kernel=kernel, optimizer=None,alpha=alpha) # disabling kernel parameters optimization + # Fit the model + gpr.fit(X, y) + # Predict mean and standard deviation + y_pred_mean, y_pred_std = gpr.predict(state_action_space, return_std=True) + + return y_pred_mean, y_pred_std + + + +# Main function for p-KRVI policy +def pi_krvi_policy(M, T, state_space, action_space,state_action_space, optimal_V,beta): + S, A, H, P, r = M + + # Initialize wandb + wandb.run.summary["episode length"] = H + wandb.run.summary["episode number"] = T + wandb.run.summary["UCB coef"] = beta + +# These arrays store the observations over all episodes + all_states=[] + all_actions=[] + all_rewards=[] + + + Qt_mean = np.zeros((H, len(state_action_space))) + Qt_std = np.zeros((H, len(state_action_space))) + Qt_estimate = np.zeros((H+1, len(state_space), len(action_space))) + for episode in range(T): + print('episode',episode) + + # Initialize arrays to store observations for the current episode + episode_states = [] + episode_actions = [] + episode_rewards = [] + + initial_state_index = np.random.randint(len(state_space)) # Sample state index + + state = state_space[initial_state_index] # Initial state + + # Collecting (s,a) pairs and (r+V(s')) from previous episodes + if episode>0: + for h in reversed(range(H)): + X_states = np.concatenate([all_states[i][h] for i in range(episode)]) + X_actions = np.concatenate([all_actions[i][h] for i in range(episode)]) + + # Reshape X_states and X_actions to be 2D arrays + X_states = X_states.reshape(-1, 1) + X_actions = X_actions.reshape(-1, 1) + + # Concatenate X_states and X_actions along axis 1 + X = np.concatenate((X_states, X_actions), axis=1) + + + Qnext = [] + for i in range(episode): + if h < H - 1: + next_state_index = np.argmin(np.abs(state_space - all_states[i][h + 1])) + Qnext.append(np.max(Qt_estimate[h + 1][next_state_index, :])) + else: + # Handle the case where h is at the last step of the episode + Qnext.append(0) # Set Qnext to 0 if we are at the last step + + y = np.array([all_rewards[i][h] + Qnext[i] for i in range(len(Qnext))]) + # if you would like to test with Standard RBF kernel, replace GP_regression_invariant with GP_regression_with_RBF + Qt_mean[h], Qt_std[h] = GP_regression_invariant(X, y, state_action_space) + Qt_mean_reshaped = Qt_mean[h].reshape((len(state_space), len(action_space))) + Qt_std_reshaped = Qt_std[h].reshape((len(state_space), len(action_space))) + + # Calculate Qt_estimate[h] from Qt_mean and Qt_std + Qt_estimate[h] = Qt_mean_reshaped + beta * Qt_std_reshaped + + + # Loop to choose actions greedily and collect observations + for h in range(H): + # Choose action greedily based on Q-values + state_index = np.argmin(np.abs(state_space - state)) # Find index of the current state + q_values=Qt_estimate[h][state_index] # Get optimistic Q-values for current state + action_index = np.argmax(q_values) # Find index of action with highest Q-value + action = action_space[action_index] # Choose action with highest Q-value + # Proceed with the chosen action + next_state = transition_dynamics(state, action,P,state_space,action_space) + reward = r[state_index, action_index] + + # Store observations + episode_states.append(state) + episode_actions.append(action) + episode_rewards.append(reward) + state = next_state #replace by next_state + + + # value function of the initial state + + episode_cum_rewards = np.sum(episode_rewards) # Simple approximation (negative of cumulative rewards) + # Compute regret for the episode + episode_regret = optimal_V[initial_state_index]- episode_cum_rewards # regret= V* - V(pi) + + metrics={"Episode_number":episode,"Episode_Regret": episode_regret, "Episode_Rewards": episode_cum_rewards} + wandb.log(metrics) + + all_states.append(np.array(episode_states)) + + all_actions.append(np.array(episode_actions)) + + all_rewards.append(np.array(episode_rewards)) + + + plt.plot(optimal_V) + plt.ylabel("optimal value function") + wandb.log({"Value iteration":wandb.Image(plt)}) + + + diff --git a/test_randomized_env.py b/test_randomized_env.py new file mode 100644 index 0000000..c11de1d --- /dev/null +++ b/test_randomized_env.py @@ -0,0 +1,109 @@ +from typing import Any, Callable, Final, Sequence, Tuple + +import numpy as np +import random + +import gymnasium as gym +from gymnasium import spaces +from gymnasium.core import ActType, ObsType, WrapperObsType + +from utils import generate_random_map + +from typing import Any, Tuple +import numpy as np +import gymnasium as gym +from gymnasium import spaces +from utils import generate_random_map, generate_random_map_with_fixed_lakes + +class FrozenLake2DStateWrapper(gym.ObservationWrapper): + def __init__(self, env: gym.Env, rescale=False, size=4, p=0.8): + super().__init__(env) + self.rescale = rescale + self.size = size + self.p = p # Probability of frozen tiles + + self.char_map = {b'S': 0, b'F': 1, b'H': 2, b'G': 3} + + self.desc = None # Will be initialized in reset + self.grid_size = size + self.start_position = None + self.goal_position = None + self.hole_positions = None + + self._initialize_map() # Initialize the first map + + def _initialize_map(self): #generate a random map where S and G are still restricted to opposite corners, holes position and number are random + """Generates a new random map and updates attributes.""" + rotate= random.randint(0, 3) + # print('rotate',rotate) + self.desc = np.array(generate_random_map_with_fixed_lakes(size=self.size, n_hole=4, rotate=rotate), dtype='c') + self.grid_size = len(self.desc) + + self.hole_positions = [] + for i, row in enumerate(self.desc): + for j, item in enumerate(row): + if item == b'H': + self.hole_positions.append([i, j]) + elif item == b'S': + self.start_position = np.array([i, j]) + elif item == b'G': + self.goal_position = np.array([i, j]) + + self.hole_positions = np.array(self.hole_positions) + # print('self.start_position',self.start_position) + # print('self.goal_position',self.goal_position) + # print('self.hole_positions',self.hole_positions) + + # Update observation space + n_holes = len(self.hole_positions) + self.observation_space = spaces.Box( + low=-0.5 * self.grid_size + 0.5, + high=3 - self.grid_size * 0.5 + 0.5, + shape=(n_holes + 2, 2), + dtype=np.float32 + ) + # ✅ Reinitialize the environment with the new map + self.env = gym.make("FrozenLake-v1", desc=self.desc.tolist(), is_slippery=False) + + def reset(self, **kwargs) -> Tuple[Any, dict]: + """Resets the environment and generates a new random map.""" + self._initialize_map() # Generate a new random map at each reset + obs, info = self.env.reset(**kwargs) + return self.observation(obs), info + + def observation(self, observation: Any) -> Any: + """Transforms the environment's observation into the 2D state representation.""" + current_pos = np.array(list(self._get_position_from_obs(observation))) + pos_obs = np.concatenate( + [ + current_pos.reshape(1, -1), + self.goal_position.reshape(1, -1), + self.hole_positions + ] + ) * 1.0 + if self.rescale: + pos_obs = pos_obs - 0.5 * self.grid_size + 0.5 + return pos_obs + + def _get_position_from_obs(self, obs: int) -> Tuple[int, int]: + """Converts observation (integer state) to grid coordinates.""" + return divmod(obs, self.grid_size) + + + + + +if __name__ == "__main__": + base_env = gym.make("FrozenLake-v1", is_slippery=False) + print(base_env.unwrapped.desc) + + # Step 2: Wrap it with FrozenLake2DStateWrapper + env = FrozenLake2DStateWrapper(base_env, rescale=True) + env.reset() + + print(env.desc) + for i in range(100): + action = int(input()) + obs, reward, terminated, truncated, info = env.step(action) + print(obs, reward, terminated, truncated, info) + diff --git a/test_rotated_reflected.py b/test_rotated_reflected.py new file mode 100644 index 0000000..e9f44ba --- /dev/null +++ b/test_rotated_reflected.py @@ -0,0 +1,136 @@ +from typing import Any, Callable, Final, Sequence, Tuple + +import numpy as np +import random + +import gymnasium as gym +from gymnasium import spaces +from gymnasium.core import ActType, ObsType, WrapperObsType + + +from typing import Any, Tuple +import numpy as np +import gymnasium as gym +from gymnasium import spaces + +class FrozenLake2DStateWrapper(gym.ObservationWrapper): + def __init__(self, env: gym.Env, rescale=False, size=4, p=0.8): + super().__init__(env) + self.rescale = rescale + self.size = size + self.p = p # Probability of frozen tiles + + self.char_map = {b'S': 0, b'F': 1, b'H': 2, b'G': 3} + + self.original_desc = env.unwrapped.desc # Will be initialized in reset + self.desc=None + self.grid_size = size + self.start_position = None + self.goal_position = None + self.hole_positions = None + + self._initialize_map() # Initialize the first map + + def _initialize_map(self): #generate a random map where S and G are still restricted to opposite corners, holes position and number are random + """Generates a new random map and updates attributes.""" + + + transform_idx = random.randint(0, 7) # 0 to 7 (8 transformations) + + if transform_idx < 4: + # Rotate by 0, 90, 180, or 270 degrees + self.desc = self.rotate_map(self.original_desc, transform_idx) + # print(f'Rotated by {transform_idx * 90} degrees') + + else: + # Reflect Y after rotation + rotation_idx = transform_idx - 4 + self.desc = self.reflect_map_y(self.rotate_map(self.original_desc, rotation_idx)) + # print(f'Rotated by {rotation_idx * 90} degrees and then reflected across Y-axis') + + # print('Transformed desc:\n', self.desc) + self.grid_size = len(self.desc) + + self.hole_positions = [] + for i, row in enumerate(self.desc): + for j, item in enumerate(row): + if item == b'H': + self.hole_positions.append([i, j]) + elif item == b'S': + self.start_position = np.array([i, j]) + elif item == b'G': + self.goal_position = np.array([i, j]) + + self.hole_positions = np.array(self.hole_positions) + # print('self.start_position',self.start_position) + # print('self.goal_position',self.goal_position) + # print('self.hole_positions',self.hole_positions) + + # Update observation space + n_holes = len(self.hole_positions) + self.observation_space = spaces.Box( + low=-0.5 * self.grid_size + 0.5, + high=3 - self.grid_size * 0.5 + 0.5, + shape=(n_holes + 2, 2), + dtype=np.float32 + ) + # ✅ Reinitialize the environment with the new map + self.env = gym.make("FrozenLake-v1", desc=self.desc.tolist(), is_slippery=False) + + def reset(self, **kwargs) -> Tuple[Any, dict]: + """Resets the environment and generates a new random map.""" + self._initialize_map() # Generate a new random map at each reset + + obs, info = self.env.reset(**kwargs) + return self.observation(obs), info + + def observation(self, observation: Any) -> Any: + """Transforms the environment's observation into the 2D state representation.""" + current_pos = np.array(list(self._get_position_from_obs(observation))) + pos_obs = np.concatenate( + [ + current_pos.reshape(1, -1), + self.goal_position.reshape(1, -1), + self.hole_positions + ] + ) * 1.0 + if self.rescale: + pos_obs = pos_obs - 0.5 * self.grid_size + 0.5 + return pos_obs + + def _get_position_from_obs(self, obs: int) -> Tuple[int, int]: + """Converts observation (integer state) to grid coordinates.""" + return divmod(obs, self.grid_size) + + def rotate_map(self, desc, rotation_idx): + """Applies a 90-degree rotation to the map `rotation_idx` times.""" + rotated_desc = np.array(desc) + for _ in range(rotation_idx): + rotated_desc = np.array(list(zip(*rotated_desc[::-1]))) + return rotated_desc + def reflect_map_x(self, desc): + """Reflects the map along the X-axis (flipping rows).""" + return np.flipud(desc) + + def reflect_map_y(self, desc): + """Reflects the map along the Y-axis (flipping columns).""" + return np.fliplr(desc) + + + + + + +if __name__ == "__main__": + base_env = gym.make("FrozenLake-v1", map_name="4x4", is_slippery=False) + # print(base_env.unwrapped.desc) + + # Step 2: Wrap it with FrozenLake2DStateWrapper + env = FrozenLake2DStateWrapper(base_env, rescale=True) + env.reset() + + for i in range(100): + action = int(input()) + obs, reward, terminated, truncated, info = env.step(action) + # print(obs, reward, terminated, truncated, info) + diff --git a/utils.py b/utils.py index 0a14a5a..17d5089 100644 --- a/utils.py +++ b/utils.py @@ -4,7 +4,7 @@ import numpy as np -from KRVI.group_rep import GroupRep +# from KRVI.group_rep import GroupRep def is_valid(board: List[List[str]], max_size: int) -> bool: @@ -81,15 +81,15 @@ def generate_random_map_with_fixed_lakes(size: int = 8, n_hole: int = 1, rotate: return desc -def construct_90deg_block_rot_groups(dim_space: int): - assert dim_space % 2 == 0 - rotation_matrix = np.array([[0, -1], - [1, 0]]) +# def construct_90deg_block_rot_groups(dim_space: int): +# assert dim_space % 2 == 0 +# rotation_matrix = np.array([[0, -1], +# [1, 0]]) - blocks = [rotation_matrix]*int(dim_space / 2) - block_diagonal_matrix = block_diag(*blocks) +# blocks = [rotation_matrix]*int(dim_space / 2) +# block_diagonal_matrix = block_diag(*blocks) - rot_group_2d = [np.eye(2), rotation_matrix, rotation_matrix @ rotation_matrix, rotation_matrix @ rotation_matrix @ rotation_matrix] - rot_group_nd = [np.eye(dim_space), block_diagonal_matrix, block_diagonal_matrix @ block_diagonal_matrix, block_diagonal_matrix @ block_diagonal_matrix @ block_diagonal_matrix] +# rot_group_2d = [np.eye(2), rotation_matrix, rotation_matrix @ rotation_matrix, rotation_matrix @ rotation_matrix @ rotation_matrix] +# rot_group_nd = [np.eye(dim_space), block_diagonal_matrix, block_diagonal_matrix @ block_diagonal_matrix, block_diagonal_matrix @ block_diagonal_matrix @ block_diagonal_matrix] - return GroupRep(rot_group_nd), GroupRep(rot_group_2d) \ No newline at end of file +# return GroupRep(rot_group_nd), GroupRep(rot_group_2d)