device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
class Learner(nn.Module):
    def __init__(self, env, hidden_dim = 256, random_prob=0.0):
        super().__init__()
        self.fc1 = nn.Linear(np.array(env.observation_space.shape).prod(), hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc_out = nn.Linear(hidden_dim, np.prod(env.action_space.shape))

        self.env = env
        self.random_prob = random_prob

    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        out = F.tanh(self.fc_out(x))
        return out

    def get_action(self, obs):
        if np.random.random() < self.random_prob:
            return self.env.action_space.sample()
        action = self.forward(torch.tensor([obs]).float().to(device))
        return np.array(action[0].detach().cpu())

%load_ext autoreload
%autoreload 2
import bc

bc_learner = Learner(env)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
bc_learner.to(device)
checkpoint_path = get_checkpoint_path("bc")
reseed(2)
bc.train(bc_learner, torch.tensor(train_obs).to(device), torch.tensor(train_actions).to(device), torch.tensor(validation_obs).to(device), torch.tensor(validation_actions).to(device), checkpoint_path, num_epochs = 1500)

import gym
import tqdm
from tqdm import tqdm
import torch.nn as nn
import torch
import matplotlib.pyplot as plt
import numpy as np
import torch.nn.functional as F
from torch.utils.data import TensorDataset, DataLoader
from torch.optim import optimizer

def train(learner, observations, actions, validation_obs, validation_acts, checkpoint_path, num_epochs=100):
    """Train function for learning a new policy using BC.
    
    Parameters:
        learner (Learner)
            A Learner object (policy)
        observations (list of numpy.ndarray)
            A list of numpy arrays of shape (7166, 11, ) 
        actions (list of numpy.ndarray)
            A list of numpy arrays of shape (7166, 3, )
        checkpoint_path (str)
            The path to save the best performing checkpoint
        num_epochs (int)
            Number of epochs to run the train function for
    
    Returns:
        learner (Learner)
            A Learner object (policy)
    """
    best_loss = float('inf')
    best_model_state = None

    loss_fn = torch.nn.MSELoss()
    optimizer = torch.optim.Adam(learner.parameters(), lr=3e-4)
    dataset = TensorDataset(observations, actions) # Create your dataset
    dataloader = DataLoader(dataset, batch_size=256, shuffle=True) # Create your dataloader

    # TODO: Complete the training loop here ###
    loss = float('inf')
    validation_losses = []
    for epoch in tqdm(range(num_epochs)):
        epoch_loss = 0
        # TODO: Iterate through dataloader and train

        # TODO: Calculate Validation Loss and Append to validation_losses (remember to use "with torch.no_grad():")


        # Saving model state if current loss is less than best loss
        if loss < best_loss:
            best_loss = loss
            best_model_state = learner.state_dict()

    # Save the best performing checkpoint
    if checkpoint_path:
        torch.save(best_model_state, checkpoint_path)
        # Save graph if not running DAGGER:
        plt.plot(np.arange(0, num_epochs), [t.cpu().numpy() for t in validation_losses])
        plt.title("Validation Loss vs. Iteration")
        plt.xlabel("Iteration")
        plt.ylabel("Validation Loss")
        plt.savefig("BC_validation_loss.png")
    
    return learner

do the TODOS