TD_ActorCritic.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tue May 28 21:16:17 2024       \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 535.171.04             Driver Version: 535.171.04   CUDA Version: 12.2     |\n",
      "|-----------------------------------------+----------------------+----------------------+\n",
      "| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |\n",
      "|                                         |                      |               MIG M. |\n",
      "|=========================================+======================+======================|\n",
      "|   0  NVIDIA GeForce RTX 3080 ...    Off | 00000000:01:00.0  On |                  N/A |\n",
      "| N/A   58C    P5              19W / 150W |     97MiB / 16384MiB |     28%      Default |\n",
      "|                                         |                      |                  N/A |\n",
      "+-----------------------------------------+----------------------+----------------------+\n",
      "                                                                                         \n",
      "+---------------------------------------------------------------------------------------+\n",
      "| Processes:                                                                            |\n",
      "|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |\n",
      "|        ID   ID                                                             Usage      |\n",
      "|=======================================================================================|\n",
      "|    0   N/A  N/A      2525      G   /usr/lib/xorg/Xorg                           93MiB |\n",
      "+---------------------------------------------------------------------------------------+\n",
      "Setting environment variable to use GPU rendering:\n",
      "env: MUJOCO_GL=egl\n",
      "Checking that the installation succeeded:\n",
      "Installation successful.\n"
     ]
    }
   ],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import grad, jit, vmap\n",
    "from jax.nn import relu, softmax\n",
    "import mujoco\n",
    "from mujoco import mjx\n",
    "import mediapy as media\n",
    "from MujocoSim import MujocoSim\n",
    "import numpy as np\n",
    "import csv\n",
    "import os\n",
    "\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax.numpy.linalg import norm\n",
    "from jax import grad, jit, vmap\n",
    "from jax.nn import relu, softmax\n",
    "import mujoco\n",
    "from mujoco import mjx\n",
    "import mediapy as media\n",
    "import mujoco.mjx\n",
    "import subprocess\n",
    "import distutils.util\n",
    "\n",
    "os.environ[\"XLA_PYTHON_CLIENT_PREALLOCATE\"] = \"false\" \n",
    "if subprocess.run('nvidia-smi').returncode:\n",
    "  raise RuntimeError(\n",
    "      'Cannot communicate with GPU. '\n",
    "      'Make sure you are using a GPU Colab runtime. '\n",
    "      'Go to the Runtime menu and select Choose runtime type.')\n",
    "\n",
    "# Add an ICD config so that glvnd can pick up the Nvidia EGL driver.\n",
    "# This is usually installed as part of an Nvidia driver package, but the Colab\n",
    "# kernel doesn't install its driver via APT, and as a result the ICD is missing.\n",
    "# (https://github.com/NVIDIA/libglvnd/blob/master/src/EGL/icd_enumeration.md)\n",
    "NVIDIA_ICD_CONFIG_PATH = '/usr/share/glvnd/egl_vendor.d/10_nvidia.json'\n",
    "if not os.path.exists(NVIDIA_ICD_CONFIG_PATH):\n",
    "  with open(NVIDIA_ICD_CONFIG_PATH, 'w') as f:\n",
    "    f.write(\"\"\"{\n",
    "    \"file_format_version\" : \"1.0.0\",\n",
    "    \"ICD\" : {\n",
    "        \"library_path\" : \"libEGL_nvidia.so.0\"\n",
    "    }\n",
    "}\n",
    "\"\"\")\n",
    "\n",
    "# Tell XLA to use Triton GEMM, this improves steps/sec by ~30% on some GPUs\n",
    "xla_flags = os.environ.get('XLA_FLAGS', '')\n",
    "xla_flags += ' --xla_gpu_triton_gemm_any=True'\n",
    "os.environ['XLA_FLAGS'] = xla_flags\n",
    "\n",
    "# Configure MuJoCo to use the EGL rendering backend (requires GPU)\n",
    "print('Setting environment variable to use GPU rendering:')\n",
    "%env MUJOCO_GL=egl\n",
    "\n",
    "try:\n",
    "  print('Checking that the installation succeeded:')\n",
    "  import mujoco\n",
    "  mujoco.MjModel.from_xml_string('<mujoco/>')\n",
    "except Exception as e:\n",
    "  raise e from RuntimeError(\n",
    "      'Something went wrong during installation. Check the shell output above '\n",
    "      'for more information.\\n'\n",
    "      'If using a hosted Colab runtime, make sure you enable GPU acceleration '\n",
    "      'by going to the Runtime menu and selecting \"Choose runtime type\".')\n",
    "\n",
    "print('Installation successful.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MujocoSim:\n",
    "\n",
    "    def __init__(self):\n",
    "        # Initialize your environment\n",
    "        # Make model, data, and renderer\n",
    "        \n",
    "        # states\n",
    "        \n",
    "        # 0-5: robot arm joints pos\n",
    "        # 6: finger pos\n",
    "        # 7-9: box pos\n",
    "        # 10-13: box quat\n",
    "        # 14-20: robot arm joints vel\n",
    "        # 21: finger vel\n",
    "        # 22-24: box vel\n",
    "        # 25-27: box angular vel\n",
    "\n",
    "        # actions\n",
    "        # 0-5 :robot arm joints\n",
    "        # 6: finger torque\n",
    "\n",
    "        self.mj_model = mujoco.MjModel.from_xml_path('simple_arm/scene.xml')\n",
    "        \n",
    "        self.mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n",
    "        self.mj_model.opt.iterations = 6\n",
    "        self.mj_model.opt.ls_iterations = 6\n",
    "        self.mj_data = mujoco.MjData(self.mj_model)\n",
    "        # renderer = mujoco.Renderer(mj_model)\n",
    "        # weight_load_target_dist_reward = 1\n",
    "        # weight_tip_to_load_position_reward = 1\n",
    "        # weight_tip_to_load_velocity_reward = 1\n",
    "        # weight_current_torque_cost= 1\n",
    "        # weight_peak_torque_cost= 1\n",
    "        # weight_timestep = 1\n",
    "        self.weights=jnp.array([-1,-45,0,1,1,1]).transpose()\n",
    "        self.load_dest=jnp.array([1,1,1]).transpose()\n",
    "\n",
    "        self.max_allowable_distance=4\n",
    "        self.max_allowable_target_error=0.1\n",
    "        # self.peak_torque=0\n",
    "\n",
    "        self.mjx_model = mjx.put_model(self.mj_model)\n",
    "        self.mjx_data = mjx.put_data(self.mj_model, self.mj_data)\n",
    "        self.mjx_data.replace(qpos=jnp.array([16e-6, -0.0007421, -0.047, 0.06 ,-4e-05 , 2.33e-5, 0.0009 ,0 ,3, 0 ,0.0198922, 1 ,0, 0 ,0]))\n",
    "\n",
    "\n",
    "        # self.p=jnp.zeros([3,6]) #TODO \n",
    "        # self.J=jnp.zeros([3,6])\n",
    "    def reset(self):\n",
    "        # Reset the environment to the initial state\n",
    "        self.mj_model = mujoco.MjModel.from_xml_path('simple_arm/scene.xml')\n",
    "        self.mj_model.opt.solver = mujoco.mjtSolver.mjSOL_CG\n",
    "        self.mj_model.opt.iterations = 6\n",
    "        self.mj_model.opt.ls_iterations = 6\n",
    "        self.mj_data = mujoco.MjData(self.mj_model)\n",
    "        # renderer = mujoco.Renderer(mj_model)\n",
    "\n",
    "        self.weights=jnp.array([1,1,1,1,1,0.1]).transpose()\n",
    "        self.load_dest=jnp.array([1,1,1]).transpose()\n",
    "\n",
    "\n",
    "\n",
    "        self.mjx_model = mjx.put_model(self.mj_model)\n",
    "        self.mjx_data = mjx.put_data(self.mj_model, self.mj_data)\n",
    "        self.mjx_data.replace(qpos=jnp.array([16e-6, -0.0007421, -0.047, 0.06 ,-4e-05 , 2.33e-5, 0.0009 ,0 ,5, 0 ,0.0198922, 1 ,0, 0 ,0]))\n",
    "    \n",
    "\n",
    "        \n",
    "        return self.get_state(self.mjx_data)\n",
    "    \n",
    "    # @jax.vmap\n",
    "    def step(self, model, data, action):\n",
    "        # Execute the action and return the new state and reward\n",
    "        # rng = jax.random.PRNGKey(0)\n",
    "        # rng = jax.random.split(rng,1024)\n",
    "\n",
    "\n",
    "        # fun=lambda rng: mjx_data.replace(ctrl=jax.random.uniform(rng, (8,)))\n",
    "        # fun_vmapped = jax.vmap(fun)\n",
    "        # batch=fun_vmapped(rng)\n",
    "        input=jnp.zeros(14)\n",
    "        input[0:8]=action\n",
    "        data.replace(ctrl=input)\n",
    "        for i in range(10):\n",
    "            mjx.step(model, data)\n",
    "        data=mjx.step(model, data)\n",
    "\n",
    "        state=self.get_state(data)\n",
    "        reward= self.get_reward(state,action)\n",
    "        # self.peak_torque=jnp.max(jnp.array([self.peak_torque, norm(jnp.array([action[0:6]]))**2]))\n",
    "        return state, reward, data\n",
    "           \n",
    "        # fun_vmapped = jax.vmap(step, in_axes=(None,0,0))\n",
    "        # batch=fun_vmapped(mjx_model, batch,rng)\n",
    "        # jit_step = jax.jit(jax.vmap(step, in_axes=(None, 0,0)))\n",
    "        # batch = jit_step(mjx_model, batch,rng)\n",
    "\n",
    "\n",
    "    def batch_step(self, states, actions):\n",
    "        batch = self.step(self.mjx_model, self.mjx_data, actions)\n",
    "    \n",
    "    def get_reward(self, state, action):\n",
    "        #update peak torque(\n",
    "        \n",
    "        return sum([self.weights[0]*norm(state[7:10]-self.load_dest), \n",
    "                    self.weights[1]*norm(state[7:10]-self.mjx_data.geom_xpos[16]),\n",
    "                    self.weights[2]*norm(self.mjx_data.geom_xpos[16]-state[22:25]) ,\n",
    "                    self.weights[3]*norm(action[0:6])**2 ,\n",
    "                    self.weights[4]*1])\n",
    "        \n",
    "    def get_state(self,data):\n",
    "\n",
    "        state=jnp.concatenate([data.qpos[0:7],data.qpos[7:14],data.qvel[0:7], data.qpos[7:13]])\n",
    "        return state\n",
    "    def isnt_done(self,state):\n",
    "        rb = jnp.array(state[7:10]).transpose()\n",
    "        rd=self.load_dest\n",
    "        rm=jnp.array([0,0,0]).transpose()\n",
    "\n",
    "        a=max([self.max_allowable_distance-norm(rb-rm),0]) \n",
    "            \n",
    "        b=max([norm(rb-rd)-self.max_allowable_target_error,0])\n",
    "        return a*b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Neural Networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "def clip_grads(grads, max_norm):\n",
    "    norm = jnp.sqrt(sum(jnp.sum(g ** 2) for g in grads.values()))\n",
    "    clip_coef = jnp.minimum(1.0, max_norm / (norm + 1e-6))\n",
    "    return {k: v * clip_coef for k, v in grads.items()}\n",
    "# Network definitions\n",
    "@jit\n",
    "def actor_network(params, state):\n",
    "\n",
    "    # print(f\"state {state}, w1 {params['W1']}, b1 {params['b1']}\")\n",
    "    \n",
    "    hidden = relu(jnp.dot(state, params['W1']) + params['b1'])\n",
    "    logits = jnp.dot(hidden, params['W2']) + params['b2']\n",
    "    return softmax(logits)\n",
    "\n",
    "@jit\n",
    "def critic_network(params, state):\n",
    "\n",
    "    hidden = relu(jnp.dot(state, params['W1']) + params['b1'])\n",
    "    value = jnp.dot(hidden, params['W2']) + params['b2']\n",
    "    return value\n",
    "\n",
    "\n",
    "def initialize_params(input_dim, hidden_dim, output_dim):\n",
    "    params = {\n",
    "        'W1': jnp.array(np.random.randn(input_dim, hidden_dim) * 0.01),\n",
    "        'b1': jnp.zeros(hidden_dim),\n",
    "        'W2': jnp.array(np.random.randn(hidden_dim, output_dim) * 0.01),\n",
    "        'b2': jnp.zeros(output_dim)\n",
    "    }\n",
    "    return params\n",
    "\n",
    "input_dim = 27  # Example input dimension\n",
    "hidden_dim = 128\n",
    "output_dim_actor = 7  # Number of actions\n",
    "output_dim_critic = 1  # Single value output\n",
    "\n",
    "actor_params = initialize_params(input_dim, hidden_dim, output_dim_actor)\n",
    "critic_params = initialize_params(input_dim, hidden_dim, output_dim_critic)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Actor Critic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ActorCritic:\n",
    "    def __init__(self, env : MujocoSim, actor_params : dict, critic_params : dict, lr : float=0.01):\n",
    "        self.env = env\n",
    "        self.actor_params = actor_params\n",
    "        self.critic_params = critic_params\n",
    "        self.lr = lr\n",
    "\n",
    "    # @jit\n",
    "    def select_action(self, state,actor_params):\n",
    "        return actor_network(actor_params, state)\n",
    "        # return jnp.ones([7,1])\n",
    "    # @jit\n",
    "    def update(self, state, action, reward, next_state, actor_params, critic_params):\n",
    "        gamma = 0.99  # Discount factor\n",
    "        max_grad_norm = 1.0  # Adjust as needed\n",
    "\n",
    "        # Compute TD target\n",
    "        value = critic_network(critic_params, state)\n",
    "        next_value = critic_network(critic_params, next_state)\n",
    "        td_target = reward + gamma * next_value\n",
    "        td_error = td_target - value\n",
    "\n",
    "\n",
    "        # print(\"tderr\", td_error)\n",
    "        # print(\"value\", value)\n",
    "        # print(\"reward\",reward)\n",
    "        # Update critic\n",
    "        def critic_loss(params):\n",
    "            value = critic_network(params, state)\n",
    "            return jnp.mean((td_target - value) ** 2)\n",
    "\n",
    "        # tde=(reward + 0.9 * critic_network(critic_params, next_state) - critic_network(critic_params, state)[0])\n",
    "\n",
    "\n",
    "        critic_grads = grad(critic_loss)(critic_params)\n",
    "        critic_grads = clip_grads(critic_grads, max_grad_norm)\n",
    "\n",
    "        critic_params = {k: critic_params[k] + self.lr * critic_grads[k] for k in critic_params}\n",
    "\n",
    "\n",
    "        key = jax.random.PRNGKey(0)\n",
    "        std_devs=jnp.ones([1,7])*50\n",
    "\n",
    "\n",
    "\n",
    "        def actor_loss(params):\n",
    "              \n",
    "            pi=jnp.abs(jax.random.normal(key,(1,7))*std_devs+actor_network(actor_params, state)[0])\n",
    "\n",
    "            pi=jnp.where(pi>0.01, pi, 0.01)\n",
    "            log_pi = jnp.log(pi)\n",
    "\n",
    "            # print(\"logpi\",log_pi)\n",
    "            return jnp.mean(log_pi * td_error)\n",
    "      \n",
    "        actor_grads = grad(actor_loss)(actor_params)\n",
    "        actor_grads = clip_grads(actor_grads, max_grad_norm)\n",
    "\n",
    "        # print(\"actorgrad\",actor_grads)\n",
    "        # print(\"acc\", actor_grads)\n",
    "        # for k in actor_params:\n",
    "        #     print(k)\n",
    "        #     for i in actor_params[k]:\n",
    "        #         print(i)\n",
    "        pi = actor_network(actor_params, state)   \n",
    "\n",
    "        actor_params = {k: actor_params[k]  + self.lr * actor_grads[k] for k in actor_params}\n",
    "        # actor_params['b2']=actor_params['b2']*pi\n",
    "        # pi = actor_network(actor_params, state)   \n",
    "        \n",
    "        # print(f\"pi: {pi}, log_pi{ jnp.log(jnp.where(pi != 0., pi, 0.0001))},\")\n",
    "        return actor_params, critic_params\n",
    "\n",
    "\n",
    "    def batch_train(self, episodes=10000, batch_size=4096):\n",
    "        \n",
    "        select_action=jax.vmap(self.select_action, in_axes=(0,None))\n",
    "        \n",
    "   \n",
    "        update=jax.jit(jax.vmap(self.update, in_axes=(0,0,0,0,None,None)))\n",
    "        # update_vmap=self.update\n",
    "  \n",
    "\n",
    "        batch=jax.vmap(lambda rng: self.env.mjx_data.replace(ctrl=jax.random.uniform(rng, (8,))))(jax.random.split(jax.random.PRNGKey(0),batch_size))\n",
    "        get_state=jax.vmap(self.env.get_state)\n",
    "        step=jax.jit(jax.vmap(self.env.step, in_axes=(None,0,0)))\n",
    "        # step=jax.vmap(self.env.step, in_axes=(None,0,0))\n",
    "        for episode in range(episodes):\n",
    "            state = get_state(batch)\n",
    "            #done = false\n",
    "            action = select_action(state, self.actor_params)\n",
    "            next_state, reward, batch = step(self.env.mjx_model, batch, action)\n",
    "            if episode==0:\n",
    "              self.log_header(jnp.mean(state, axis=0), jnp.mean(action, axis=0), jnp.mean(reward, axis=0))\n",
    "            \n",
    "            #visualize\n",
    "            # renderer = mujoco.Renderer(self.env.mj_model)\n",
    "            # frames=[]\n",
    "            # framerate=100\n",
    "            print(\"ep\", episode)\n",
    "            for i in range(32000):\n",
    "                # print(i)\n",
    "                action = select_action(state, self.actor_params)\n",
    "\n",
    "                next_state, reward, batch = step(self.env.mjx_model, batch, action)\n",
    "                \n",
    "                # jax.debug.print(f\"next_state.shape {next_state.shape} mean {jnp.mean(next_state, axis=0)}\")\n",
    "\n",
    "\n",
    "                actor_params,  critic_params = update(state, action, reward, next_state, self.actor_params, self.critic_params)\n",
    "                for key in actor_params:\n",
    "                    self.actor_params[key]=jnp.mean(actor_params[key], axis=0)\n",
    "                    self.critic_params[key]=jnp.mean(critic_params[key], axis=0)\n",
    "                # for i in range(batch_size):    \n",
    "                #      self.actor_params,  self.critic_params = update(state[i], action[i], reward[i], next_state[i], self.actor_params, self.critic_params)\n",
    "                #     for key in actor_params:\n",
    "                #         self.actor_params[key]=actor_params[key]\n",
    "\n",
    "                #     for key in critic_params:\n",
    "                #         self.critic_params[key]=critic_params[key]\n",
    "                state = next_state\n",
    "                \n",
    "    \n",
    "    \n",
    "\n",
    "                # self.log_line(jnp.mean(c, axis=0), jnp.mean(action, axis=0), jnp.mean(reward, axis=0))\n",
    "                # batched_mj_data = mjx.get_data(self.env.mj_model, batch)    \n",
    "                # mj_data=batched_mj_data[0]\n",
    "                # renderer.update_scene(mj_data)\n",
    "                # pixels = renderer.render()\n",
    "                # frames.append(pixels)\n",
    "            # media.show_video(frames, fps=framerate)\n",
    "                \n",
    "                # batched_mj_data = mjx.get_data(self.env.mj_model, batch)\n",
    "\n",
    "    def train(self, episodes=1):\n",
    "   \n",
    "        reset=jax.jit(self.env.reset())\n",
    "\n",
    "        for episode in range(episodes):\n",
    "            reset()\n",
    "            batch=self.env.mjx_data.replace(ctrl=jnp.ones([8]))\n",
    "            \n",
    "            state = jnp.zeros(27)\n",
    "\n",
    "            done = False\n",
    "            action = self.select_action(state, self.actor_params)\n",
    "            next_state, reward, batch = self.env.step(self.env.mjx_model, batch, action)\n",
    "            if episode==0:\n",
    "              self.log_header(state,action,reward)\n",
    "            renderer = mujoco.Renderer(self.env.mj_model)\n",
    "\n",
    "\n",
    "            frames=[]\n",
    "            framerate=100\n",
    "            print(\"ep\", episode)\n",
    "            for i in range(10):\n",
    "\n",
    "                action = self.select_action(state, self.actor_params)\n",
    "\n",
    "                \n",
    "                next_state, reward, batch = self.env.step(self.env.mjx_model, batch, action)\n",
    "\n",
    "                # jax.debug.print(f\"next_state.shape {next_state.shape} mean {jnp.mean(next_state, axis=0)}\")\n",
    "                for i in range(1):    \n",
    "                    actor_params, critic_params = self.update(state, action, reward, next_state, self.actor_params, self.critic_params)\n",
    "                    for key in actor_params:\n",
    "                        self.actor_params[key]=actor_params[key]\n",
    "\n",
    "                    for key in critic_params:\n",
    "                        self.critic_params[key]=critic_params[key]\n",
    "                state = next_state\n",
    "                \n",
    "    \n",
    "\n",
    "                print(f\"action: {action}\")\n",
    "                self.log_line(state,action,reward)\n",
    "                \n",
    "\n",
    "    # def update_parameters(self, mm, critic_params):\n",
    "\n",
    "    def log_header(self,state, action, reward):\n",
    "        header_text=[]\n",
    "        for i in range(len(state)):\n",
    "            header_text.append(\"state\"+\"_\"+str(i) )\n",
    "        for i in range(len(action)):\n",
    "            header_text.append(\"action\"+\"_\"+str(i))\n",
    "        header_text.append(\"reward\")\n",
    "        filename=\"ac_log.csv\"\n",
    "\n",
    "        \n",
    "        \n",
    "        user_input = 'y'\n",
    "        if user_input == 'y': \n",
    "            file = 'logs/'+filename\n",
    "            try:\n",
    "                os.remove (file)\n",
    "                data_f = open('logs/'+filename, 'a',newline='')\n",
    "            except FileNotFoundError:\n",
    "                data_f = open('logs/'+filename, 'x',newline='')\n",
    "            # data_f = open('../opy_logs/'+filename, 'a',newline='')\n",
    "            self.data_writer = csv.writer(data_f)\n",
    "            Headers = header_text\n",
    "            print(Headers)\n",
    "            self.data_writer.writerow(Headers) \n",
    "        else:\n",
    "            data_f = open('logs/'+filename, 'a',newline='')\n",
    "            data_writer = csv.writer(data_f)\n",
    "    def log_newline(self):    \n",
    "        self.log_text=self.log_text+\"0\\n\"        \n",
    "    def log_line(self, state, action, reward):\n",
    "\n",
    "        LogList=[]\n",
    "        # print(f\"next_state.shape {state.shape} mean {str(state)}\")\n",
    "        # LogList.append(state)\n",
    "        for i in state:\n",
    "            LogList.append(str(i)) \n",
    "        for i in action:\n",
    "            LogList.append(str(i)) \n",
    "        LogList.append(reward) \n",
    "        self.data_writer.writerow(LogList) \n",
    "    def save_parameters(self):\n",
    "        LogList=[]\n",
    "        # print(f\"next_state.shape {state.shape} mean {str(state)}\")\n",
    "        \n",
    "        filename=\"params.csv\"\n",
    "        file = 'logs/'+filename\n",
    "        try:\n",
    "            os.remove (file)\n",
    "            data_f = open('logs/'+filename, 'a',newline='')\n",
    "        except FileNotFoundError:\n",
    "            data_f = open('logs/'+filename, 'x',newline='')\n",
    "        # LogList.append(state)\n",
    "        data_writer = csv.writer(data_f)\n",
    "        for key in self.actor_params:\n",
    "            LogList.append(key)\n",
    "            for item in self.actor_params[key]:\n",
    "                LogList.append(str(item))\n",
    "        # print(f\"next_state.shape {state.shape} mean {str(state)}\")\n",
    "        # LogList.append(state)\n",
    "        for key in self.critic_params:\n",
    "            LogList.append(key)\n",
    "            for item in self.actor_params[key]:\n",
    "                LogList.append(str(item))\n",
    "        \n",
    "        data_writer.writerow(LogList)\n",
    "        # Vectorized operations using vmap\n",
    "        # v_update = vmap(self.update, in_axes=(0, 0, 0, 0))\n",
    "        # v_update(states, actions, rewards, next_states)\n",
    "        \n",
    "\n",
    "    # @jit\n",
    "    # def batch_select_action(self, states):\n",
    "    #     # Vectorized action selection\n",
    "    #     v_select = vmap(self.select_action)\n",
    "    #     return self.v_se lect(states)\n",
    "    def use_parameters(self):\n",
    "        self.env.reset()\n",
    "        if(1):\n",
    "            batch=self.env.mjx_data.replace(ctrl=jnp.ones([8]))\n",
    "            \n",
    "            state = jnp.zeros(27)\n",
    "\n",
    "            done = False\n",
    "            action = self.select_action(state, self.actor_params)\n",
    "            next_state, reward, batch = self.env.step(self.env.mjx_model, batch, action)\n",
    "            if 1:\n",
    "              self.log_header(state,action,reward)\n",
    "            renderer = mujoco.Renderer(self.env.mj_model)\n",
    "\n",
    "\n",
    "            frames=[]\n",
    "            framerate=100\n",
    "            step=jax.jit(self.env.step)\n",
    "            print(\"use_parameters\", 1)\n",
    "            for i in range(100):\n",
    "\n",
    "                action = self.select_action(state, self.actor_params)\n",
    "\n",
    "                \n",
    "                next_state, reward, batch =step(self.env.mjx_model, batch, action)\n",
    "\n",
    "                # jax.debug.print(f\"next_state.shape {next_state.shape} mean {jnp.mean(next_state, axis=0)}\")\n",
    "                # for i in range(1):    \n",
    "                #     actor_params, critic_params = self.update(state, action, reward, next_state, self.actor_params, self.critic_params)\n",
    "                #     for key in actor_params:\n",
    "                #         self.actor_params[key]=actor_params[key]\n",
    "\n",
    "                #     for key in critic_params:\n",
    "                #         self.critic_params[key]=critic_params[key]\n",
    "                state = next_state\n",
    "                # self.log_line(jnp.mean(state, axis=0), jnp.mean(action, axis=0), jnp.mean(reward, axis=0))\n",
    "                mj_data = mjx.get_data(self.env.mj_model, batch)    \n",
    "                # mj_data=batched_mj_data[0]\n",
    "                renderer.update_scene(mj_data)\n",
    "                pixels = renderer.render()\n",
    "                frames.append(pixels)\n",
    "                self.log_line(state,action,reward)               \n",
    "            media.show_video(frames, fps=framerate)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Running"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "begin\n",
      "['state_0', 'state_1', 'state_2', 'state_3', 'state_4', 'state_5', 'state_6', 'state_7', 'state_8', 'state_9', 'state_10', 'state_11', 'state_12', 'state_13', 'state_14', 'state_15', 'state_16', 'state_17', 'state_18', 'state_19', 'state_20', 'state_21', 'state_22', 'state_23', 'state_24', 'state_25', 'state_26', 'action_0', 'action_1', 'action_2', 'action_3', 'action_4', 'action_5', 'action_6', 'reward']\n",
      "ep 0\n",
      "ep 1\n",
      "ep 2\n",
      "ep 3\n",
      "ep 4\n",
      "ep 5\n",
      "ep 6\n",
      "ep 7\n",
      "ep 8\n",
      "ep 9\n",
      "ep 10\n",
      "ep 11\n",
      "ep 12\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 9\u001b[0m\n\u001b[1;32m      7\u001b[0m \u001b[38;5;66;03m# # Train the model\u001b[39;00m\n\u001b[1;32m      8\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mbegin\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m----> 9\u001b[0m \u001b[43mac\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mbatch_train\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     10\u001b[0m \u001b[38;5;66;03m# batch_train=jax.jit(ac.batch_train)\u001b[39;00m\n\u001b[1;32m     11\u001b[0m \u001b[38;5;66;03m# batch_train()\u001b[39;00m\n\u001b[1;32m     16\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msave params\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "Cell \u001b[0;32mIn[4], line 105\u001b[0m, in \u001b[0;36mActorCritic.batch_train\u001b[0;34m(self, episodes, batch_size)\u001b[0m\n\u001b[1;32m    101\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;241m32000\u001b[39m):\n\u001b[1;32m    102\u001b[0m     \u001b[38;5;66;03m# print(i)\u001b[39;00m\n\u001b[1;32m    103\u001b[0m     action \u001b[38;5;241m=\u001b[39m select_action(state, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mactor_params)\n\u001b[0;32m--> 105\u001b[0m     next_state, reward, batch \u001b[38;5;241m=\u001b[39m \u001b[43mstep\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43menv\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmjx_model\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbatch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43maction\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    107\u001b[0m     \u001b[38;5;66;03m# jax.debug.print(f\"next_state.shape {next_state.shape} mean {jnp.mean(next_state, axis=0)}\")\u001b[39;00m\n\u001b[1;32m    110\u001b[0m     actor_params,  critic_params \u001b[38;5;241m=\u001b[39m update(state, action, reward, next_state, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mactor_params, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mcritic_params)\n",
      "File \u001b[0;32m~/.local/lib/python3.12/site-packages/mujoco/mjx/_src/dataclasses.py:81\u001b[0m, in \u001b[0;36mdataclass.<locals>.clz_from_iterable\u001b[0;34m(meta, data)\u001b[0m\n\u001b[1;32m     78\u001b[0m   meta \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mtuple\u001b[39m(to_meta(f, x) \u001b[38;5;28;01mfor\u001b[39;00m f \u001b[38;5;129;01min\u001b[39;00m meta_fields)\n\u001b[1;32m     79\u001b[0m   \u001b[38;5;28;01mreturn\u001b[39;00m data, meta\n\u001b[0;32m---> 81\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mclz_from_iterable\u001b[39m(meta, data):\n\u001b[1;32m     83\u001b[0m   \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mfrom_meta\u001b[39m(field, meta):\n\u001b[1;32m     84\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m field\u001b[38;5;241m.\u001b[39mtype \u001b[38;5;129;01mis\u001b[39;00m np\u001b[38;5;241m.\u001b[39mndarray:\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "\n",
    "# Instantiate your environment\n",
    "env = MujocoSim()\n",
    "\n",
    "# Create ActorCritic instance\n",
    "ac = ActorCritic(env, actor_params, critic_params)\n",
    "\n",
    "# # Train the model\n",
    "print(\"begin\")\n",
    "ac.batch_train()\n",
    "# batch_train=jax.jit(ac.batch_train)\n",
    "# batch_train()\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "print(\"save params\")\n",
    "ac.save_parameters()\n",
    "print(\"visualize one run\")\n",
    "ac.use_parameters()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "save params\n"
     ]
    },
    {
     "ename": "NameError",
     "evalue": "name 'ac' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[1], line 3\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[38;5;66;03m# ac.use_parameters()\u001b[39;00m\n\u001b[1;32m      2\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msave params\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m----> 3\u001b[0m \u001b[43mac\u001b[49m\u001b[38;5;241m.\u001b[39msave_parameters()\n\u001b[1;32m      4\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mvisualize one run\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m      5\u001b[0m ac\u001b[38;5;241m.\u001b[39muse_parameters()\n",
      "\u001b[0;31mNameError\u001b[0m: name 'ac' is not defined"
     ]
    }
   ],
   "source": [
    "# ac.use_parameters()\n",
    "# print(\"save params\")\n",
    "# ac.save_parameters()\n",
    "# print(\"visualize one run\")\n",
    "# ac.use_parameters()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}