From 875df231eaf44e9b57526bfe2d87896d191a9b2b Mon Sep 17 00:00:00 2001 From: ayakayal Date: Sun, 11 May 2025 14:15:35 +0100 Subject: [PATCH 1/6] clean synthetic and frozen lake exp --- KRVI_algo_test_rotated_invariant_fulloptim.py | 468 ++++++++++++++++++ ...o_test_rotated_invariant_fulloptim_eval.py | 449 +++++++++++++++++ common_framework2.py | 38 ++ framework.py | 341 +++++++++++++ invariant_kernel_fixed.py | 223 +++++++++ script.py | 69 +++ test_kernelized_RL_invariant.py | 216 ++++++++ test_randomized_env.py | 115 +++++ test_rotated_reflected.py | 155 ++++++ utils.py | 95 ++++ 10 files changed, 2169 insertions(+) create mode 100644 KRVI_algo_test_rotated_invariant_fulloptim.py create mode 100644 KRVI_algo_test_rotated_invariant_fulloptim_eval.py create mode 100644 common_framework2.py create mode 100644 framework.py create mode 100644 invariant_kernel_fixed.py create mode 100644 script.py create mode 100644 test_kernelized_RL_invariant.py create mode 100644 test_randomized_env.py create mode 100644 test_rotated_reflected.py create mode 100644 utils.py diff --git a/KRVI_algo_test_rotated_invariant_fulloptim.py b/KRVI_algo_test_rotated_invariant_fulloptim.py new file mode 100644 index 0000000..31101d9 --- /dev/null +++ b/KRVI_algo_test_rotated_invariant_fulloptim.py @@ -0,0 +1,468 @@ +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() + # print('initial_state',initial_state) + #print('self.current_rotation',self.env.current_rotation) + #print('self.desc',self.env.desc) + + 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): #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 +# """ +# # Ensure inputs are torch tensors and use double precision +# 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), likelihood=GaussianLikelihood(noise_constraint=GreaterThan(1e-4))) #,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 = True +# #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 = True + +# # 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": 5, "raw_samples":512}) +# del mll + +# except botorch.exceptions.errors.ModelFittingError as e: +# print("Model fitting failed. Printing parameters:") +# for param_name, param in model.named_parameters(): +# print(f'Parameter name: {param_name:42} value: {param.item()} requires_grad: {param.requires_grad}') +# # Optionally, print the traceback for debugging +# traceback.print_exc() +# # print('X',X) +# # print('Y',Y) + +# # **Memory Cleanup** +# del X, y # Safe to delete +# torch.cuda.empty_cache() # Free GPU memory + + +# return model + + 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 + """ + # Ensure inputs are torch tensors and use double precision + 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="trial2", 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..14dffcb --- /dev/null +++ b/KRVI_algo_test_rotated_invariant_fulloptim_eval.py @@ -0,0 +1,449 @@ +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 test_rotated_env 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() + # print('initial_state',initial_state) + #print('self.current_rotation',self.env.current_rotation) + #print('self.desc',self.env.desc) + + 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 + """ + # Ensure inputs are torch tensors and use double precision + 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_rewards = [] + + 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) + + #episode_reward.append(reward) + total_reward += reward + + if done or truncated: + break + + state = next_state + + test_returns.append(total_reward) + #test_rewards.append(sum(episode_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="EEE_IQL", 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=5000, 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/common_framework2.py b/common_framework2.py new file mode 100644 index 0000000..15d63ba --- /dev/null +++ b/common_framework2.py @@ -0,0 +1,38 @@ +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 + + # Create subdirectory for P_kernel if it doesn't exist + 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) + + # Save P, r, and optimal_value_function + 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..623cf59 --- /dev/null +++ b/framework.py @@ -0,0 +1,341 @@ +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) + #K += np.eye(X.shape[0]) * 1e-6 + + #K += self.kernel(X_transformed, Y) / len(self.group) + #print('k',K.shape) + # Add jitter only if K is square + # if K.shape[0] == K.shape[1]: + # print('yesss jitter') + # K += np.eye(K.shape[0]) * 1e-6 # Small jitter for stability + + + #K += np.eye(X.shape[0]) * 1e-6 + #K = (K + K.T) / 2 + + + return K + + def diag(self, X): + return np.diag(self.__call__(X)) + # Compute the diagonal entries of the group-invariant kernel + # X = np.atleast_2d(X) + # K_diag = np.zeros(X.shape[0]) + + # # Average the kernel diagonals over the transformations + # for g in self.group: + # X_transformed = np.array([g(x) for x in X]) + # K_diag += np.diag(self.kernel(X_transformed, X)) / len(self.group) + + # return K_diag + + 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 + +# # # Define transformation groups for state and action + +# state_transformations = [lambda x: x, flip] +# action_transformations = [lambda x: x, flip] + + +# # # Combine state and action transformations for 2D inputs (S x A) +# group_SA = [ +# lambda x: np.array([g1(x[0]), g2(x[1])]) +# for g1, g2 in product(state_transformations, action_transformations) +# ] +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 +] + + +# # Combine state, action, and next state transformations for 3D inputs (S x A x S') +# group_SASp = [ +# lambda x: np.array([g1(x[0]), g2(x[1]), g3(x[2])]) +# for g1, g2, g3 in product(state_transformations, action_transformations, state_transformations) +# ] + +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 = np.random.randn(X.shape[0], 1) + #y = np.random.rand(X.shape[0]) # Random values (normal distribution) + y = gp.sample_y(X, 1) + #print('X',X.shape) + #print('y',y) + # Fit the Gaussian Process Regressor + #print('alpha',alpha) + gpr = GaussianProcessRegressor(kernel=kernel, optimizer=None, alpha=alpha) + K = gpr.kernel(X,X) + eigvals = np.linalg.eigvalsh(K) + #print("Kernel matrix eigenvalues:", eigvals) + + gpr.fit(X, y) + #print('after fitting') + + # Generate all possible input points for prediction (across state-action space) + #print('length state space',len(state_space)) + 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) + #print('y_pred',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))) + #print('r',r) + + 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) + + # # Scatter plot of predicted rewards + # state_indices = np.arange(len(state_space)) + # action_indices = np.arange(len(action_space)) + # for state_idx in state_indices: + # for action_idx in action_indices: + # x = state_space[state_idx] + # y = action_space[action_idx] + # z = r[state_idx, action_idx] + # ax.scatter(x, y, z, color='red', s=50) + + 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): + + # Choose state and action indices + # state_indices = [0, 50, 99] # Indices of states in state_space + # action_indices = [0, 50, 99] # Indices of actions in action_space + + 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..44f6d65 --- /dev/null +++ b/invariant_kernel_fixed.py @@ -0,0 +1,223 @@ +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 +#from botorch.models import SingleTaskGP +#from botorch.fit import fit_gpytorch_mll + + +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/script.py b/script.py new file mode 100644 index 0000000..32ae4e4 --- /dev/null +++ b/script.py @@ -0,0 +1,69 @@ +import test_kernelized_RL_invariant +#from framework import value_iteration_episodic, transition_P_RKHS, reward_RKHS, plot_reward_gp_3d, plot_transition_probabilities, InvariantKernel +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.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') # Example format: 2024-11-19_15-30-45 + 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) + # print('P',P) + # print('P.shape',P.shape) + # print('r',r) + # print('r.shape',r.shape) + # print('optimal_value_function',optimal_value_function) + # print('optimal_value_function_shape',optimal_value_function.shape) + + + + # 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.001 # UCB coefficient + NUM_RUNS=10 + #wandb.init(project="kernelized_RL") + + + # pi_krvi_policy(M, T,state_space, action_space,optimal_value_function,beta) + + #,name=f"run_{run}" + for run in range(NUM_RUNS): + # Start a new run for each iteration + wandb.init(project="IQL", 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..6af5842 --- /dev/null +++ b/test_kernelized_RL_invariant.py @@ -0,0 +1,216 @@ +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": + #print('length_scale',l) + 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) # We should not evaluate at discrete state-action space, return gpr + + 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) #correct + 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 + +# # Function for GP regression to estimate Q-values using Matérn kernel +# def GP_regression_with_Matern(X, y, state_action_space, alpha=1.0, nu=1.5): +# # Define the Matérn kernel +# kernel = Matern(length_scale=1.0, length_scale_bounds="fixed", nu=nu) +# wandb.run.summary["kernel type"] = kernel +# # Create GPR model +# gpr = GaussianProcessRegressor(kernel=kernel, alpha=alpha, optimizer=None) +# # 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 + #initial_state_index=2 + ###print('initial state index',initial_state_index) + state = state_space[initial_state_index] # Initial state + + # Collecting (s,a) pairs and (r+V(s')) from previous episodes + if episode>0: + ## complete the code + for h in reversed(range(H)): + ###print('h',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) + ###print('concatinated states',X_states) + X_actions = X_actions.reshape(-1, 1) + ###print('concatinated actions',X_actions) + + # Concatenate X_states and X_actions along axis 1 + X = np.concatenate((X_states, X_actions), axis=1) + #print('X',X) + ###print('concatinated rewards',np.concatenate([np.array(all_rewards[i][h]) for i in range(episode)])) + + Qnext = [] + for i in range(episode): + if h < H - 1: + next_state_index = np.argmin(np.abs(state_space - all_states[i][h + 1])) + ###print('next_state_index',next_state_index) + 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 + ###print('Qnext',Qnext) + # Concatenate rewards and Qnext along axis 0 and add them element-wise + #print('all_rewards',all_rewards) + #print('Qnext',Qnext) + y = np.array([all_rewards[i][h] + Qnext[i] for i in range(len(Qnext))]) + #print('y',y) + #y = np.concatenate([np.array(all_rewards[i][h]) + Qnext[i] for i in range(len(Qnext))]) + ###print('y',y) + Qt_mean[h], Qt_std[h] = GP_regression_with_RBF(X, y, state_action_space) + ###print('Qt_mean',Qt_mean[h]) + ###print('Qt_std',Qt_std[h]) + 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 + ###print('Qt_estimate[h]',Qt_estimate[h]) + + + # 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 + ###print('state_index',state_index) + q_values=Qt_estimate[h][state_index] # Get optimistic Q-values for current state + ###print('q_values',q_values) + 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 + ###print('action',action) + # 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_index = np.argmin(np.abs(state_space - next_state)) # Find closest next state index + state = next_state #replace by next_state + + + # value function of the initial state + + #print('Optimal V value iteration',optimal_value_function[initial_state_index]) + 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) + #print('episode_regret',episode_regret) + #print('episode cum rewards',episode_cum_rewards) + metrics={"Episode_number":episode,"Episode_Regret": episode_regret, "Episode_Rewards": episode_cum_rewards} + wandb.log(metrics) + #cum_rewards.append( episode_cum_rewards) + ##print('episode_states',episode_states) + all_states.append(np.array(episode_states)) + #print('all_states',all_states) + ##print('episode_actions',episode_actions) + all_actions.append(np.array(episode_actions)) + #print('all_actions',all_actions) + ##print('episode_rewards',episode_rewards) + all_rewards.append(np.array(episode_rewards)) + ##print('all_rewards',all_rewards) + # At the end of the function, inside the loop where episodes are being iterated + + # if episode == T - 1: # Check if it's the last episode + # # Convert Qt_estimate to a table format + # table_data = [] + # for h in range(H): + # for i, state in enumerate(state_space): + # for j, action in enumerate(action_space): + # table_data.append([h, state, action, Qt_estimate[h][i][j]]) + # # Log the table + #wandb.log({"Qt_estimate_table": wandb.Table(data=table_data, columns=["Step", "State", "Action", "Q_estimate"])}) + + 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..1661ab3 --- /dev/null +++ b/test_randomized_env.py @@ -0,0 +1,115 @@ +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 # Ensure this is correctly imported + +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(size=self.size, p=self.p, rotate=random.randint(0, 3)), dtype='c') + 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 = FrozenLake2DStateWrapper(None, rescale=True) # Start with no env, let wrapper create it + env.reset() + + # env = gym.make('FrozenLake-v1', map_name='4x4', is_slippery=False) + # print(env.unwrapped.desc) + # env = FrozenLake2DStateWrapper(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..c559d05 --- /dev/null +++ b/test_rotated_reflected.py @@ -0,0 +1,155 @@ +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_type = random.choice(["rotate", "reflect_x", "reflect_y"]) # Randomly pick a transformation + # transform_value = random.randint(0, 3) if transform_type == "rotate" else None + + # if transform_type == "rotate": + # print('Rotating by', transform_value * 90, 'degrees') + # self.desc = self.rotate_map(self.original_desc, transform_value) + # elif transform_type == "reflect_x": + # print('Reflecting across X-axis') + # self.desc = self.reflect_map_x(self.original_desc) + # elif transform_type == "reflect_y": + # print('Reflecting across Y-axis') + # self.desc = self.reflect_map_y(self.original_desc) + + # print('Transformed desc:\n', self.desc) + 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 = FrozenLake2DStateWrapper(None, rescale=True) # Start with no env, let wrapper create it + env.reset() + + # env = gym.make('FrozenLake-v1', map_name='4x4', is_slippery=False) + # print(env.unwrapped.desc) + # env = FrozenLake2DStateWrapper(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/utils.py b/utils.py new file mode 100644 index 0000000..0c74a1c --- /dev/null +++ b/utils.py @@ -0,0 +1,95 @@ +# DFS to check that it's a valid path. +from typing import List, Optional +from scipy.linalg import block_diag + +import numpy as np + +# from KRVI.group_rep import GroupRep + + +def is_valid(board: List[List[str]], max_size: int) -> bool: + frontier, discovered = [], set() + frontier.append((0, 0)) + while frontier: + r, c = frontier.pop() + if not (r, c) in discovered: + discovered.add((r, c)) + directions = [(1, 0), (0, 1), (-1, 0), (0, -1)] + for x, y in directions: + r_new = r + x + c_new = c + y + if r_new < 0 or r_new >= max_size or c_new < 0 or c_new >= max_size: + continue + if board[r_new][c_new] == "G": + return True + if board[r_new][c_new] != "H": + frontier.append((r_new, c_new)) + return False + + +def generate_random_map(size: int = 8, p: float = 0.8, rotate:int = 0) -> List[str]: + """Generates a random valid map (one that has a path from start to goal) + + Args: + size: size of each side of the grid + p: probability that a tile is frozen + + Returns: + A random valid map + """ + valid = False + board = [] # initialize to make pyright happy + + while not valid: + p = min(1, p) + board = np.random.choice(["F", "H"], (size, size), p=[p, 1 - p]) + board[0][0] = "S" + board[-1][-1] = "G" + valid = is_valid(board, size) + desc = ["".join(x) for x in board] + + for i in range(rotate): + desc = list(zip(*desc[::-1])) + return desc + +def generate_random_map_with_fixed_lakes(size: int = 8, n_hole: int = 1, rotate: int = 0) -> List[str]: + """Generates a random valid map (one that has a path from start to goal) + + Args: + size: size of each side of the grid + n_hole: number of lakes to place + + Returns: + A random valid map, or the empty list [] is none exists + """ + valid = False + + while not valid: + board = np.full((size, size), "F") # initialize to make pyright happy + board[0][0] = "S" + board[-1][-1] = "G" + positions = list(range(size**2 - 2)) # number of options for holes + hole_positions = np.random.choice(positions, size=n_hole, replace = False) + for position in hole_positions: + pos_tup = divmod(position + 1, size) # cannot be 0, 0 + board[pos_tup] = "H" + valid = is_valid(board, size) + desc = ["".join(x) for x in board] + + for i in range(rotate): + desc = list(zip(*desc[::-1])) + return desc + + +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) + + 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 From 753843e08e3d41c496506d79d84e92981f31c070 Mon Sep 17 00:00:00 2001 From: ayakayal Date: Sun, 11 May 2025 14:56:20 +0100 Subject: [PATCH 2/6] Add requirements.txt with project dependencies --- KRVI_algo_test_rotated_invariant_fulloptim.py | 2 +- ...o_test_rotated_invariant_fulloptim_eval.py | 2 +- requirements.txt | 77 +++++++++++++++++++ script.py | 8 +- 4 files changed, 83 insertions(+), 6 deletions(-) create mode 100644 requirements.txt diff --git a/KRVI_algo_test_rotated_invariant_fulloptim.py b/KRVI_algo_test_rotated_invariant_fulloptim.py index 31101d9..3c7a128 100644 --- a/KRVI_algo_test_rotated_invariant_fulloptim.py +++ b/KRVI_algo_test_rotated_invariant_fulloptim.py @@ -400,7 +400,7 @@ def predict_with_gp(self, model, states_batch, actions_batch): 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="trial2", help="wandb project name") #IQL_project_invariant + 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") diff --git a/KRVI_algo_test_rotated_invariant_fulloptim_eval.py b/KRVI_algo_test_rotated_invariant_fulloptim_eval.py index 14dffcb..adcc067 100644 --- a/KRVI_algo_test_rotated_invariant_fulloptim_eval.py +++ b/KRVI_algo_test_rotated_invariant_fulloptim_eval.py @@ -381,7 +381,7 @@ def evaluate(self, Qt, action_space, num_episodes=40): #evaluating on 40 differe 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="EEE_IQL", help="wandb project name") #IQL_project_invariant + 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=5000, help="Number of training iterations (T)") parser.add_argument("--seed", type=int, default=0, help="random seed") 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 index 32ae4e4..950f7c2 100644 --- a/script.py +++ b/script.py @@ -19,7 +19,7 @@ def main(): H = 10 # Length of each episode P_kernel= 'RBF' - alpha=0.001 + alpha=0.01 #0.001 # Define save directory #save_dir = f'Rewards_and_P_invariant_NEW' #subdir = os.path.join(save_dir, P_kernel) @@ -48,8 +48,8 @@ def main(): # Run experiments with p-KRVI policy T = 1000 # Number of episodes - beta=0.001 # UCB coefficient - NUM_RUNS=10 + beta=0.1 # UCB coefficient (0.001) + NUM_RUNS=20 #wandb.init(project="kernelized_RL") @@ -58,7 +58,7 @@ def main(): #,name=f"run_{run}" for run in range(NUM_RUNS): # Start a new run for each iteration - wandb.init(project="IQL", reinit=True) + 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 From 80ed804c571ba652d6ab7e0ac8ce436df7b5147b Mon Sep 17 00:00:00 2001 From: ayakayal Date: Sun, 11 May 2025 15:06:02 +0100 Subject: [PATCH 3/6] add readme --- README.md | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/README.md b/README.md index 72161e4..b559ad7 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,24 @@ # IQL Kernelized Q-learning with invariances. + +## Synthetic Experiments +To run synthetic experiments: + +```bash +python 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 +``` + +## 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 +``` From 238b9da4b7cf00eec06e812d083591f17c8b9a58 Mon Sep 17 00:00:00 2001 From: ayakayal Date: Thu, 15 May 2025 10:27:53 +0100 Subject: [PATCH 4/6] Final clean up Frozen Lake experiments --- KRVI_algo_test_rotated_invariant_fulloptim.py | 71 +------------------ ...o_test_rotated_invariant_fulloptim_eval.py | 13 +--- invariant_kernel_fixed.py | 3 +- test_randomized_env.py | 8 +-- test_rotated_reflected.py | 19 ----- utils.py | 18 ++--- 6 files changed, 16 insertions(+), 116 deletions(-) diff --git a/KRVI_algo_test_rotated_invariant_fulloptim.py b/KRVI_algo_test_rotated_invariant_fulloptim.py index 3c7a128..6ebfe0a 100644 --- a/KRVI_algo_test_rotated_invariant_fulloptim.py +++ b/KRVI_algo_test_rotated_invariant_fulloptim.py @@ -188,10 +188,7 @@ def train(self, T: int): episode_rewards = [] initial_state, info = self.env.reset() - # print('initial_state',initial_state) - #print('self.current_rotation',self.env.current_rotation) - #print('self.desc',self.env.desc) - + state= preprocess_state(initial_state) @@ -249,71 +246,8 @@ def train(self, T: int): # 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 -# """ -# # Ensure inputs are torch tensors and use double precision -# 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), likelihood=GaussianLikelihood(noise_constraint=GreaterThan(1e-4))) #,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 = True -# #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 = True - -# # 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": 5, "raw_samples":512}) -# del mll - -# except botorch.exceptions.errors.ModelFittingError as e: -# print("Model fitting failed. Printing parameters:") -# for param_name, param in model.named_parameters(): -# print(f'Parameter name: {param_name:42} value: {param.item()} requires_grad: {param.requires_grad}') -# # Optionally, print the traceback for debugging -# traceback.print_exc() -# # print('X',X) -# # print('Y',Y) - -# # **Memory Cleanup** -# del X, y # Safe to delete -# torch.cuda.empty_cache() # Free GPU memory - - -# return model - def GP_regression_torch(self, X, y): #I removed normalization + def GP_regression_torch(self, X, y): """ Gaussian Process regression using PyTorch. @@ -321,7 +255,6 @@ def GP_regression_torch(self, X, y): #I removed normalization :param y: Target tensor of shape (n_samples,) :return: Trained GP model """ - # Ensure inputs are torch tensors and use double precision X = torch.tensor(X, dtype=torch.float32,device=device) assert not torch.isnan(X).any() y = torch.tensor(y, dtype=torch.float32,device=device) diff --git a/KRVI_algo_test_rotated_invariant_fulloptim_eval.py b/KRVI_algo_test_rotated_invariant_fulloptim_eval.py index adcc067..a55c955 100644 --- a/KRVI_algo_test_rotated_invariant_fulloptim_eval.py +++ b/KRVI_algo_test_rotated_invariant_fulloptim_eval.py @@ -27,7 +27,6 @@ print('device',device) from test_randomized_env import FrozenLake2DStateWrapper #from test_rotated_reflected import FrozenLake2DStateWrapper -#from test_rotated_env import FrozenLake2DStateWrapper from invariant_kernel_fixed import InvariantKernel, apply_rotation_group import csv @@ -188,9 +187,7 @@ def train(self, T: int): episode_rewards = [] initial_state, info = self.env.reset() - # print('initial_state',initial_state) - #print('self.current_rotation',self.env.current_rotation) - #print('self.desc',self.env.desc) + state= preprocess_state(initial_state) @@ -263,7 +260,6 @@ def GP_regression_torch(self, X, y): #I removed normalization :param y: Target tensor of shape (n_samples,) :return: Trained GP model """ - # Ensure inputs are torch tensors and use double precision X = torch.tensor(X, dtype=torch.float32,device=device) assert not torch.isnan(X).any() y = torch.tensor(y, dtype=torch.float32,device=device) @@ -331,7 +327,6 @@ def predict_with_gp(self, model, states_batch, actions_batch): def evaluate(self, Qt, action_space, num_episodes=40): #evaluating on 40 different environments test_returns = [] - #test_rewards = [] test_env = FrozenLake2DStateWrapper(gym.make("FrozenLake-v1", is_slippery=False), rescale=True) @@ -354,7 +349,6 @@ def evaluate(self, Qt, action_space, num_episodes=40): #evaluating on 40 differe next_state, reward, done, truncated, _ = test_env.step(action) next_state = preprocess_state(next_state) - #episode_reward.append(reward) total_reward += reward if done or truncated: @@ -363,7 +357,6 @@ def evaluate(self, Qt, action_space, num_episodes=40): #evaluating on 40 differe state = next_state test_returns.append(total_reward) - #test_rewards.append(sum(episode_reward)) return np.mean(test_returns), np.std(test_returns) @@ -381,9 +374,9 @@ def evaluate(self, Qt, action_space, num_episodes=40): #evaluating on 40 differe 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("--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=5000, help="Number of training iterations (T)") + 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") diff --git a/invariant_kernel_fixed.py b/invariant_kernel_fixed.py index 44f6d65..746a536 100644 --- a/invariant_kernel_fixed.py +++ b/invariant_kernel_fixed.py @@ -14,8 +14,7 @@ import torch import numpy as np from scipy.linalg import block_diag -#from botorch.models import SingleTaskGP -#from botorch.fit import fit_gpytorch_mll + class InvariantKernel(gpytorch.kernels.Kernel): diff --git a/test_randomized_env.py b/test_randomized_env.py index 1661ab3..c11de1d 100644 --- a/test_randomized_env.py +++ b/test_randomized_env.py @@ -13,7 +13,7 @@ import numpy as np import gymnasium as gym from gymnasium import spaces -from utils import generate_random_map, generate_random_map_with_fixed_lakes # Ensure this is correctly imported +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): @@ -36,7 +36,6 @@ def _initialize_map(self): #generate a random map where S and G are still restri """Generates a new random map and updates attributes.""" rotate= random.randint(0, 3) # print('rotate',rotate) - #self.desc = np.array(generate_random_map(size=self.size, p=self.p, rotate=random.randint(0, 3)), dtype='c') 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) @@ -100,13 +99,8 @@ def _get_position_from_obs(self, obs: int) -> Tuple[int, int]: # Step 2: Wrap it with FrozenLake2DStateWrapper env = FrozenLake2DStateWrapper(base_env, rescale=True) - #env = FrozenLake2DStateWrapper(None, rescale=True) # Start with no env, let wrapper create it env.reset() - # env = gym.make('FrozenLake-v1', map_name='4x4', is_slippery=False) - # print(env.unwrapped.desc) - # env = FrozenLake2DStateWrapper(env, rescale = True) - # env.reset() print(env.desc) for i in range(100): action = int(input()) diff --git a/test_rotated_reflected.py b/test_rotated_reflected.py index c559d05..e9f44ba 100644 --- a/test_rotated_reflected.py +++ b/test_rotated_reflected.py @@ -34,20 +34,7 @@ def __init__(self, env: gym.Env, rescale=False, size=4, p=0.8): 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_type = random.choice(["rotate", "reflect_x", "reflect_y"]) # Randomly pick a transformation - # transform_value = random.randint(0, 3) if transform_type == "rotate" else None - - # if transform_type == "rotate": - # print('Rotating by', transform_value * 90, 'degrees') - # self.desc = self.rotate_map(self.original_desc, transform_value) - # elif transform_type == "reflect_x": - # print('Reflecting across X-axis') - # self.desc = self.reflect_map_x(self.original_desc) - # elif transform_type == "reflect_y": - # print('Reflecting across Y-axis') - # self.desc = self.reflect_map_y(self.original_desc) - # print('Transformed desc:\n', self.desc) transform_idx = random.randint(0, 7) # 0 to 7 (8 transformations) if transform_idx < 4: @@ -140,14 +127,8 @@ def reflect_map_y(self, desc): # Step 2: Wrap it with FrozenLake2DStateWrapper env = FrozenLake2DStateWrapper(base_env, rescale=True) - #env = FrozenLake2DStateWrapper(None, rescale=True) # Start with no env, let wrapper create it env.reset() - # env = gym.make('FrozenLake-v1', map_name='4x4', is_slippery=False) - # print(env.unwrapped.desc) - # env = FrozenLake2DStateWrapper(env, rescale = True) - # env.reset() - # print(env.desc) for i in range(100): action = int(input()) obs, reward, terminated, truncated, info = env.step(action) diff --git a/utils.py b/utils.py index 0c74a1c..f553dfb 100644 --- a/utils.py +++ b/utils.py @@ -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) \ No newline at end of file From 06fab191f53e512edb2aab14f2b6313721758fdf Mon Sep 17 00:00:00 2001 From: ayakayal Date: Thu, 15 May 2025 10:49:40 +0100 Subject: [PATCH 5/6] final clean up synthetic experiments --- common_framework2.py | 2 - framework.py | 70 ++++----------------------------- script.py | 16 +------- test_kernelized_RL_invariant.py | 70 +++++---------------------------- 4 files changed, 20 insertions(+), 138 deletions(-) diff --git a/common_framework2.py b/common_framework2.py index 15d63ba..a28a604 100644 --- a/common_framework2.py +++ b/common_framework2.py @@ -9,7 +9,6 @@ def save_rewards_and_transitions(P_kernel, alpha, save_dir): action_space = np.linspace(-1, 1, num=10).reshape(-1, 1) H = 10 - # Create subdirectory for P_kernel if it doesn't exist subdir = os.path.join(save_dir, P_kernel) os.makedirs(subdir, exist_ok=True) @@ -17,7 +16,6 @@ def save_rewards_and_transitions(P_kernel, alpha, save_dir): r = reward_RKHS(P_kernel, state_space, action_space, subdir, alpha) optimal_value_function = value_iteration_episodic(state_space, action_space, r, P, H) - # Save P, r, and optimal_value_function 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) diff --git a/framework.py b/framework.py index 623cf59..cf61557 100644 --- a/framework.py +++ b/framework.py @@ -54,34 +54,14 @@ def group_invariant_kernel(self, X, Y=None, eval_gradient=False): #print('Transformed kernel shape:', transformed_kernel.shape) K += transformed_kernel / len(self.group) - #K += np.eye(X.shape[0]) * 1e-6 - - #K += self.kernel(X_transformed, Y) / len(self.group) - #print('k',K.shape) - # Add jitter only if K is square - # if K.shape[0] == K.shape[1]: - # print('yesss jitter') - # K += np.eye(K.shape[0]) * 1e-6 # Small jitter for stability - - - #K += np.eye(X.shape[0]) * 1e-6 - #K = (K + K.T) / 2 + return K def diag(self, X): return np.diag(self.__call__(X)) - # Compute the diagonal entries of the group-invariant kernel - # X = np.atleast_2d(X) - # K_diag = np.zeros(X.shape[0]) - - # # Average the kernel diagonals over the transformations - # for g in self.group: - # X_transformed = np.array([g(x) for x in X]) - # K_diag += np.diag(self.kernel(X_transformed, X)) / len(self.group) - - # return K_diag + def is_stationary(self): return True @@ -98,28 +78,14 @@ def square(x): def flip(x): return - x -# # # Define transformation groups for state and action - -# state_transformations = [lambda x: x, flip] -# action_transformations = [lambda x: x, flip] - -# # # Combine state and action transformations for 2D inputs (S x A) -# group_SA = [ -# lambda x: np.array([g1(x[0]), g2(x[1])]) -# for g1, g2 in product(state_transformations, action_transformations) -# ] 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 ] -# # Combine state, action, and next state transformations for 3D inputs (S x A x S') -# group_SASp = [ -# lambda x: np.array([g1(x[0]), g2(x[1]), g3(x[2])]) -# for g1, g2, g3 in product(state_transformations, action_transformations, state_transformations) -# ] + group_SASp = [ lambda x: np.array([x[0],x[1],x[2]]), # Identity for both state and action @@ -147,23 +113,15 @@ def reward_RKHS(P_kernel, state_space, action_space, subdir=None, alpha=0.5): # Sample y values from the Gaussian process with the group-invariant kernel gp = GaussianProcessRegressor(kernel=kernel) - #y = np.random.randn(X.shape[0], 1) - #y = np.random.rand(X.shape[0]) # Random values (normal distribution) + y = gp.sample_y(X, 1) - #print('X',X.shape) - #print('y',y) - # Fit the Gaussian Process Regressor - #print('alpha',alpha) + gpr = GaussianProcessRegressor(kernel=kernel, optimizer=None, alpha=alpha) K = gpr.kernel(X,X) eigvals = np.linalg.eigvalsh(K) - #print("Kernel matrix eigenvalues:", eigvals) gpr.fit(X, y) - #print('after fitting') - - # Generate all possible input points for prediction (across state-action space) - #print('length state space',len(state_space)) + values = np.linspace(-1, 1, len(state_space)) all_possible_inputs = np.array(list(product(values, repeat=2))) # State-action pairs in 2D @@ -171,13 +129,11 @@ def reward_RKHS(P_kernel, state_space, action_space, subdir=None, alpha=0.5): 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) - #print('y_pred',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))) - #print('r',r) return r @@ -274,15 +230,7 @@ def plot_reward_gp_3d(r, state_space, action_space, save_dir=None): surf = ax.plot_surface(state_mesh, action_mesh, r, cmap='viridis',vmin=0,vmax=1) fig.colorbar(surf, shrink=0.5, aspect=5) - # # Scatter plot of predicted rewards - # state_indices = np.arange(len(state_space)) - # action_indices = np.arange(len(action_space)) - # for state_idx in state_indices: - # for action_idx in action_indices: - # x = state_space[state_idx] - # y = action_space[action_idx] - # z = r[state_idx, action_idx] - # ax.scatter(x, y, z, color='red', s=50) + ax.set_xlabel('s',fontsize=20) ax.set_ylabel('a',fontsize=20) @@ -302,9 +250,7 @@ def plot_reward_gp_3d(r, state_space, action_space, save_dir=None): def plot_transition_probabilities(P, state_space,action_space, save_dir=None): - # Choose state and action indices - # state_indices = [0, 50, 99] # Indices of states in state_space - # action_indices = [0, 50, 99] # Indices of actions in action_space + state_indices = [0, 5, 9] # Indices of states in state_space action_indices = [0, 5, 9] # Indices of actions in action_space diff --git a/script.py b/script.py index 950f7c2..df78fdf 100644 --- a/script.py +++ b/script.py @@ -1,5 +1,4 @@ import test_kernelized_RL_invariant -#from framework import value_iteration_episodic, transition_P_RKHS, reward_RKHS, plot_reward_gp_3d, plot_transition_probabilities, InvariantKernel from common_framework2 import save_rewards_and_transitions, load_saved_data import numpy as np import wandb @@ -24,7 +23,7 @@ def main(): #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') # Example format: 2024-11-19_15-30-45 + 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) @@ -34,14 +33,8 @@ def main(): save_rewards_and_transitions(P_kernel, alpha, save_dir) P, r, optimal_value_function = load_saved_data(subdir) - # print('P',P) - # print('P.shape',P.shape) - # print('r',r) - # print('r.shape',r.shape) - # print('optimal_value_function',optimal_value_function) - # print('optimal_value_function_shape',optimal_value_function.shape) - + # Define environment M = (S,A,H,P,r) M = (S, A, H, P, r) @@ -50,12 +43,7 @@ def main(): T = 1000 # Number of episodes beta=0.1 # UCB coefficient (0.001) NUM_RUNS=20 - #wandb.init(project="kernelized_RL") - - - # pi_krvi_policy(M, T,state_space, action_space,optimal_value_function,beta) - #,name=f"run_{run}" for run in range(NUM_RUNS): # Start a new run for each iteration wandb.init(project="IQL_submission_test", reinit=True) diff --git a/test_kernelized_RL_invariant.py b/test_kernelized_RL_invariant.py index 6af5842..47f1fba 100644 --- a/test_kernelized_RL_invariant.py +++ b/test_kernelized_RL_invariant.py @@ -24,7 +24,6 @@ def GP_regression_invariant(X, y, state_action_space,P_kernel='RBF',l=1,alpha=1e kernel = GroupInvariantKernel(base_kernel="Matern", length_scale=l, smoothness=2.5, group=group_SA) elif P_kernel == "RBF": - #print('length_scale',l) kernel = GroupInvariantKernel(base_kernel="RBF", length_scale=l, group=group_SA) # Define the RBF kernel @@ -37,7 +36,7 @@ def GP_regression_invariant(X, y, state_action_space,P_kernel='RBF',l=1,alpha=1e # 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) # We should not evaluate at discrete state-action space, return gpr + y_pred_mean, y_pred_std = gpr.predict(state_action_space, return_std=True) return y_pred_mean, y_pred_std @@ -52,7 +51,6 @@ def GP_regression_with_RBF(X, y, state_action_space,l=1,alpha=1e-10): wandb.run.summary["alpha_gp"] = alpha wandb.run.summary["length_scale"] = l # Create GPR model - #gpr = GaussianProcessRegressor(kernel=kernel, optimizer=None) #correct gpr = GaussianProcessRegressor(kernel=kernel, optimizer=None,alpha=alpha) # disabling kernel parameters optimization # Fit the model gpr.fit(X, y) @@ -61,19 +59,6 @@ def GP_regression_with_RBF(X, y, state_action_space,l=1,alpha=1e-10): return y_pred_mean, y_pred_std -# # Function for GP regression to estimate Q-values using Matérn kernel -# def GP_regression_with_Matern(X, y, state_action_space, alpha=1.0, nu=1.5): -# # Define the Matérn kernel -# kernel = Matern(length_scale=1.0, length_scale_bounds="fixed", nu=nu) -# wandb.run.summary["kernel type"] = kernel -# # Create GPR model -# gpr = GaussianProcessRegressor(kernel=kernel, alpha=alpha, optimizer=None) -# # 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 @@ -103,67 +88,49 @@ def pi_krvi_policy(M, T, state_space, action_space,state_action_space, optimal_V episode_rewards = [] initial_state_index = np.random.randint(len(state_space)) # Sample state index - #initial_state_index=2 - ###print('initial state index',initial_state_index) + state = state_space[initial_state_index] # Initial state # Collecting (s,a) pairs and (r+V(s')) from previous episodes if episode>0: - ## complete the code for h in reversed(range(H)): - ###print('h',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) - ###print('concatinated states',X_states) X_actions = X_actions.reshape(-1, 1) - ###print('concatinated actions',X_actions) # Concatenate X_states and X_actions along axis 1 X = np.concatenate((X_states, X_actions), axis=1) - #print('X',X) - ###print('concatinated rewards',np.concatenate([np.array(all_rewards[i][h]) for i in range(episode)])) + Qnext = [] for i in range(episode): if h < H - 1: next_state_index = np.argmin(np.abs(state_space - all_states[i][h + 1])) - ###print('next_state_index',next_state_index) 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 - ###print('Qnext',Qnext) - # Concatenate rewards and Qnext along axis 0 and add them element-wise - #print('all_rewards',all_rewards) - #print('Qnext',Qnext) + y = np.array([all_rewards[i][h] + Qnext[i] for i in range(len(Qnext))]) - #print('y',y) - #y = np.concatenate([np.array(all_rewards[i][h]) + Qnext[i] for i in range(len(Qnext))]) - ###print('y',y) + Qt_mean[h], Qt_std[h] = GP_regression_with_RBF(X, y, state_action_space) - ###print('Qt_mean',Qt_mean[h]) - ###print('Qt_std',Qt_std[h]) 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 - ###print('Qt_estimate[h]',Qt_estimate[h]) # 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 - ###print('state_index',state_index) q_values=Qt_estimate[h][state_index] # Get optimistic Q-values for current state - ###print('q_values',q_values) 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 - ###print('action',action) # Proceed with the chosen action next_state = transition_dynamics(state, action,P,state_space,action_space) reward = r[state_index, action_index] @@ -172,41 +139,24 @@ def pi_krvi_policy(M, T, state_space, action_space,state_action_space, optimal_V episode_states.append(state) episode_actions.append(action) episode_rewards.append(reward) - #state_index = np.argmin(np.abs(state_space - next_state)) # Find closest next state index state = next_state #replace by next_state # value function of the initial state - #print('Optimal V value iteration',optimal_value_function[initial_state_index]) 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) - #print('episode_regret',episode_regret) - #print('episode cum rewards',episode_cum_rewards) + metrics={"Episode_number":episode,"Episode_Regret": episode_regret, "Episode_Rewards": episode_cum_rewards} wandb.log(metrics) - #cum_rewards.append( episode_cum_rewards) - ##print('episode_states',episode_states) + all_states.append(np.array(episode_states)) - #print('all_states',all_states) - ##print('episode_actions',episode_actions) + all_actions.append(np.array(episode_actions)) - #print('all_actions',all_actions) - ##print('episode_rewards',episode_rewards) + all_rewards.append(np.array(episode_rewards)) - ##print('all_rewards',all_rewards) - # At the end of the function, inside the loop where episodes are being iterated - - # if episode == T - 1: # Check if it's the last episode - # # Convert Qt_estimate to a table format - # table_data = [] - # for h in range(H): - # for i, state in enumerate(state_space): - # for j, action in enumerate(action_space): - # table_data.append([h, state, action, Qt_estimate[h][i][j]]) - # # Log the table - #wandb.log({"Qt_estimate_table": wandb.Table(data=table_data, columns=["Step", "State", "Action", "Q_estimate"])}) + plt.plot(optimal_V) plt.ylabel("optimal value function") From b093691a10ad1df010fe96055e14b71fc3b37b26 Mon Sep 17 00:00:00 2001 From: ayakayal Date: Thu, 15 May 2025 15:18:26 +0100 Subject: [PATCH 6/6] final edits --- README.md | 27 +++++++++++++++++++++++++-- test_kernelized_RL_invariant.py | 4 ++-- 2 files changed, 27 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index b559ad7..6bfea7a 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@ # 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: @@ -7,14 +7,26 @@ 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: @@ -22,3 +34,14 @@ 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/test_kernelized_RL_invariant.py b/test_kernelized_RL_invariant.py index 47f1fba..7992730 100644 --- a/test_kernelized_RL_invariant.py +++ b/test_kernelized_RL_invariant.py @@ -115,8 +115,8 @@ def pi_krvi_policy(M, T, state_space, action_space,state_action_space, optimal_V 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))]) - - Qt_mean[h], Qt_std[h] = GP_regression_with_RBF(X, y, state_action_space) + # 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)))