main_file.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import gymnasium as gym\n",
    "from stable_baselines3 import SAC, TD3, A2C\n",
    "import os\n",
    "from gymnasium import RewardWrapper\n",
    "import numpy as np\n",
    "\n",
    "from gymnasium.wrappers import ClipAction\n",
    "\n",
    "from langchain.schema.messages import HumanMessage, AIMessage\n",
    "from langchain.memory import ConversationBufferMemory\n",
    "from langchain.prompts import (\n",
    "    ChatPromptTemplate,\n",
    "    MessagesPlaceholder,\n",
    "    SystemMessagePromptTemplate,\n",
    "    HumanMessagePromptTemplate,\n",
    ")\n",
    "from langchain.chains.llm import LLMChain\n",
    "from langchain.chat_models import ChatOpenAI"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Basic Function Definitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dir = \"models\"\n",
    "log_dir = \"logs\"\n",
    "os.makedirs(model_dir, exist_ok=True)\n",
    "os.makedirs(log_dir, exist_ok=True)\n",
    "\n",
    "def train(env, sb3_algo):\n",
    "    # episode is 1000 steps\n",
    "    if sb3_algo == 'SAC':\n",
    "        model = SAC('MlpPolicy', env, verbose=0, device='cuda', tensorboard_log=log_dir)\n",
    "    elif sb3_algo == 'TD3':\n",
    "        model = TD3('MlpPolicy', env, verbose=0, device='cuda', tensorboard_log=log_dir)\n",
    "    elif sb3_algo == 'A2C':\n",
    "        model = A2C('MlpPolicy', env, verbose=0, device='cuda', tensorboard_log=log_dir)\n",
    "    else:\n",
    "        print('Algorithm not found')\n",
    "        return\n",
    "\n",
    "    TIMESTEPS = 25000\n",
    "    iters = 0\n",
    "    while True:\n",
    "        iters += 1\n",
    "        model.learn(total_timesteps=TIMESTEPS, reset_num_timesteps=False)\n",
    "        model.save(f\"{model_dir}/{sb3_algo}_{TIMESTEPS*iters}\")\n",
    "def test(env, sb3_algo, path_to_model):\n",
    "    if sb3_algo == 'SAC':\n",
    "        model = SAC.load(path_to_model, env=env)\n",
    "    elif sb3_algo == 'TD3':\n",
    "        model = TD3.load(path_to_model, env=env)\n",
    "    elif sb3_algo == 'A2C':\n",
    "        model = A2C.load(path_to_model, env=env)\n",
    "    else:\n",
    "        print('Algorithm not found')\n",
    "        return\n",
    "\n",
    "    obs = env.reset()\n",
    "    done = False\n",
    "    extra_steps = 500\n",
    "    while True:\n",
    "        action, _ = model.predict(obs)\n",
    "        obs, _, done, _ = env.step(action)\n",
    "        \n",
    "        # Print average reward on that step\n",
    "        if done:\n",
    "            extra_steps -= 1\n",
    "            if extra_steps < 0:\n",
    "                break\n"
   ]
  },
 
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Environment Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gym import RewardWrapper\n",
    "import matplotlib.pyplot as plt\n",
    "env_name = 'Humanoid-v4'\n",
    "env = gym.make(env_name, render_mode='Human', exclude_current_positions_from_observation=False)\n",
    "\n",
    "step = 0\n",
    "past_rewards = []\n",
    "\n",
    "past_torso_positions = []\n",
    "past_max_torso_position = [1.4]\n",
    "def jump_reward(reward):\n",
    "    global step, past_rewards, past_max_torso_position, past_torso_positions\n",
    "    \n",
    "    \n",
    "    x_torso = env.state_vector()[0]\n",
    "    y_torso = env.state_vector()[1]\n",
    "    z_torso = env.state_vector()[2] \n",
    "    \n",
    "    past_torso_positions.append(z_torso) \n",
    "    \n",
    "    if z_torso < 0:\n",
    "        salto = z_torso*5\n",
    "    else:\n",
    "        salto = z_torso*20 \n",
    "        \n",
    "    if z_torso > past_max_torso_position[-1]:\n",
    "        past_max_torso_position.append(z_torso)\n",
    "        salto = z_torso*100\n",
    "\n",
    "    reward=reward + salto\n",
    "    past_rewards.append(reward) \n",
    "    # Every 1000 steps, print average reward\n",
    "    if step % 1000 == 0:\n",
    "        # we clear plots \n",
    "        print(f\"Max torso position on this iter: {max(past_torso_positions[-1000:])}\")\n",
    "        print(f\"Average reward: {np.mean(past_rewards[-1000:])}\")\n",
    "        print(f\"Max torso position: {past_max_torso_position[-1]}\\n\")\n",
    "        \n",
    "    step += 1\n",
    "    return reward\n",
    "\n",
    "\n",
    "def jump_reward_2(reward):\n",
    "    global step,past_max_torso_position, past_rewards, past_torso_positions\n",
    "    step += 1\n",
    "    z_torso = env.state_vector()[2]\n",
    "    past_torso_positions.append(z_torso)\n",
    "    salto = 0\n",
    "    \n",
    "    if z_torso < 0:\n",
    "        salto += z_torso*5\n",
    "    else:\n",
    "        # Avoid giving reward for holding position\n",
    "        if len(past_torso_positions) > 10:\n",
    "            if abs(past_torso_positions[-10] - z_torso) > 0.2:\n",
    "                salto += z_torso*20\n",
    "        \n",
    "    if z_torso > past_max_torso_position[-1]:\n",
    "        salto += z_torso*100\n",
    "        past_max_torso_position.append(z_torso)\n",
    "    \n",
    "    # If it holds position, then penalize \n",
    "    if len(past_torso_positions) > 100:\n",
    "        if z_torso- past_torso_positions[-100] < 0.3:\n",
    "            salto -= 20 \n",
    "            \n",
    "    if step % 1000 == 0:\n",
    "        print(f\"Max torso position on this iter: {max(past_torso_positions[-1000:])}\")\n",
    "        print(f\"Average reward: {np.mean(past_rewards[-1000:])}\")\n",
    "        print(f\"Max torso position: {past_max_torso_position[-1]}\\n\")\n",
    "        \n",
    "    reward = salto\n",
    "    past_rewards.append(reward) \n",
    "    return reward\n",
    "\n",
    "def jump_reward_3(reward):\n",
    "    global step, past_max_torso_position, past_rewards, past_torso_positions\n",
    "    step += 1\n",
    "    z_torso = env.state_vector()[2]\n",
    "    jump_reward = 0\n",
    "    \n",
    "    if len(past_torso_positions) > 10:\n",
    "        if z_torso < 1.4:  # Define a threshold for low height\n",
    "            jump_reward = -5\n",
    "            \n",
    "    # Reward for increasing height\n",
    "    height_increase = (z_torso - past_max_torso_position[-1] if past_max_torso_position else 0)\n",
    "    #height_increase = z_torso - (past_torso_positions[-1] if past_torso_positions else 0)\n",
    "    if height_increase > 0 and z_torso > 0:\n",
    "        #jump_reward = height_increase * 500  # Increase multiplier as needed\n",
    "        jump_reward = z_torso * 500  # Increase multiplier as needed\n",
    "\n",
    "    # Penalize for staying too low or static\n",
    "        #elif abs(past_torso_positions[-10] - z_torso) < 0.3:  # Not moving significantly\n",
    "            #jump_reward -= 10\n",
    "\n",
    "    # Update past torso positions and max position\n",
    "    past_torso_positions.append(z_torso)\n",
    "    if z_torso > past_max_torso_position[-1]:\n",
    "        past_max_torso_position.append(z_torso)\n",
    "        #jump_reward += height_increase*200\n",
    "\n",
    "    # Penalize for unnatural postures if necessary\n",
    "    # posture_penalty = calculate_posture_penalty(env)\n",
    "    # jump_reward -= posture_penalty\n",
    "\n",
    "    # Update reward\n",
    "    reward += jump_reward\n",
    "    past_rewards.append(reward)\n",
    "\n",
    "    # Periodic logging\n",
    "    if step % 1000 == 0:\n",
    "        print(f\"Max torso position on this iter: {max(past_torso_positions)}\")\n",
    "        past_torso_positions = []\n",
    "        print(f\"Average reward: {np.mean(past_rewards[-1000:])}\")\n",
    "        print(f\"Max torso position: {past_max_torso_position[-1]}\\n\")\n",
    "        \n",
    "\n",
    "    return reward \n",
    "     \n",
    "\n",
    "memory = {\n",
    "    \"past_z_coordinates\": [1.38],\n",
    "    \"max_z_coordinates\": [1.38],\n",
    "    \"steps\": 0,\n",
    "    \"past_rewards\": []\n",
    "    \n",
    "}    \n",
    "def jump_reward_4(reward):\n",
    "    global memory\n",
    "    reward=0\n",
    "    \n",
    "    memory[\"steps\"] += 1\n",
    "    x_torso = env.state_vector()[0]\n",
    "    y_torso = env.state_vector()[1]\n",
    "    z_torso = env.state_vector()[2]\n",
    "    diff = z_torso - memory[\"past_z_coordinates\"][-1]\n",
    "    \n",
    "    if abs(x_torso) > 0.1 or abs(y_torso) > 0.1:\n",
    "        reward = reward - 10\n",
    "    # We penalize if it stays put for too long\n",
    "    if abs(np.mean(memory[\"past_z_coordinates\"][-10:]) - z_torso) < 1:\n",
    "        reward = reward - 10 \n",
    "        \n",
    "    if diff > 0:\n",
    "        reward += diff*100\n",
    "        \n",
    "    if z_torso > memory[\"max_z_coordinates\"][-1]:\n",
    "        memory[\"max_z_coordinates\"].append(z_torso)\n",
    "        reward += diff*300\n",
    "        \n",
    "    memory[\"past_z_coordinates\"].append(z_torso)\n",
    "    \n",
    "    print(f\"x_torso: {x_torso}, y_torso: {y_torso}, z_torso: {z_torso}, diff: {diff}, reward: {reward}\", end=\"\\r\") \n",
    "    if memory[\"steps\"] % 1000 == 0:\n",
    "        print(f\"Max torso position on this iter: {max(memory['past_z_coordinates'][-1000:])}\")\n",
    "        print(f\"Average reward: {np.mean(past_rewards[-1000:])}\")\n",
    "        print(f\"Max torso position: {memory['max_z_coordinates'][-1]}\\n\")\n",
    "        print(memory[\"past_rewards\"][-1000:])\n",
    "        \n",
    "    past_rewards.append(reward) \n",
    "    return reward \n",
    "\n",
    "\n",
    "# Initialize global variables to keep track of past data\n",
    "step = 0\n",
    "past_rewards = []\n",
    "max_height_reached = -float('inf')  # Use negative infinity as starting point\n",
    "\n",
    "def jump_reward_5(reward):\n",
    "    global step, past_rewards, max_height_reached\n",
    "    \n",
    "    observation = env.state_vector() \n",
    "    # Extract observations\n",
    "    z_torso = observation[2]  # z-coordinate of the torso\n",
    "    torso_orientations = observation[1:5]  # x, y, z, w-orientation of the torso\n",
    "    \n",
    "    # Define constants for reward calculation\n",
    "    JUMP_REWARD_WEIGHT = 20\n",
    "    ORIENTATION_PENALTY_WEIGHT = 5\n",
    "    \n",
    "    # Reward for jumping\n",
    "    jump_reward = z_torso * JUMP_REWARD_WEIGHT if z_torso > 0 else 0\n",
    "    \n",
    "    # Update the maximum height reached\n",
    "    if z_torso > max_height_reached:\n",
    "        max_height_reached = z_torso\n",
    "    \n",
    "    # Penalty for not maintaining upright orientation\n",
    "    orientation_penalty = sum(abs(o) for o in torso_orientations) * ORIENTATION_PENALTY_WEIGHT\n",
    "    \n",
    "    # Calculate total reward\n",
    "    reward = jump_reward - orientation_penalty\n",
    "    \n",
    "    \n",
    "    # Track rewards\n",
    "    past_rewards.append(reward)\n",
    "    step += 1\n",
    "    \n",
    "    # Reset variables at the end of an episode (1000 steps)\n",
    "    if step % 1000 == 0:\n",
    "        print(f\"Max torso position on this iter: {max_height_reached}\")\n",
    "        print(f\"Average reward: {np.mean(past_rewards[-1000:])}\")\n",
    "        step = 0\n",
    "        past_rewards.clear()\n",
    "        max_height_reached = -float('inf')\n",
    "    \n",
    "    return reward\n",
    "\n",
    "max_height_reached = -float('inf')  # Use negative infinity as starting point\n",
    "steps = 0\n",
    "past_rewards = []\n",
    "def jump_reward_6(_):\n",
    "    global max_height_reached, steps, past_rewards\n",
    "    z_torso = env.state_vector()[2]\n",
    "    reward = -1\n",
    "    if z_torso > max_height_reached:\n",
    "        max_height_reached = z_torso\n",
    "        reward = 300\n",
    "    \n",
    "    if steps % 1000 == 0:\n",
    "        print(f\"Average reward: {np.mean(past_rewards[-1000:])}\")\n",
    "        print(f\"Max torso position: {max_height_reached}\\n\")    \n",
    "      \n",
    "    steps += 1\n",
    "    past_rewards.append(reward)\n",
    "        \n",
    "    return reward\n",
    "\n",
    "step = 0\n",
    "past_rewards = []\n",
    "past_max_x_position = [-9999]\n",
    "\n",
    "def walk_reward(reward):\n",
    "    global step, past_rewards, past_max_x_position\n",
    "    x_coor_torso = env.state_vector()[0]\n",
    "    \n",
    "    # Reward for moving forward on the x-axis\n",
    "    displacement_reward = x_coor_torso\n",
    "    \n",
    "    # Check if the current x position is the furthest the humanoid has been\n",
    "    if x_coor_torso > past_max_x_position[-1]:\n",
    "        past_max_x_position.append(x_coor_torso)\n",
    "        \n",
    "    if step%1000 == 0:\n",
    "        print(f\"Max x position on this iter: {max(past_max_x_position[-1000:])}\")\n",
    "        \n",
    "    reward += displacement_reward\n",
    "    past_rewards.append(reward)\n",
    "    step += 1\n",
    "    return reward\n",
    "    \n",
    "env = gym.wrappers.TransformReward(env, walk_reward)\n",
    "env.reset()\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"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\"\n",
    "gpt3 = ChatOpenAI(temperature=0, model_name=\"gpt-3.5-turbo-16k\")\n",
    "gpt4 = ChatOpenAI(temperature=0.7, model_name=\"gpt-4-1106-preview\")\n",
    "\n",
    "\n",
    "input_template = \"\"\"{0}\"\"\"\n",
    "system_template = \"\"\"\n",
    "Given an objective, you write the code for the reward function of a humanoid robot. Here is the observation space available to you:\n",
    "    | Num | Observation                                                                                                     | Min  | Max | Name (in corresponding XML file) | Joint | Unit                       |\n",
    "| --- | --------------------------------------------------------------------------------------------------------------- | ---- | --- | -------------------------------- | ----- | -------------------------- |\n",
    "| 0   | x-coordinate of the torso (centre)                                                                              | -Inf | Inf | root                             | free  | position (m)               |\n",
    "| 1   | y-coordinate of the torso (centre)                                                                              | -Inf | Inf | root                             | free  | position (m)               |\n",
    "| 2   | z-coordinate of the torso (centre)                                                                              | -Inf | Inf | root                             | free  | position (m)               |\n",
    "| 3   | x-orientation of the torso (centre)                                                                             | -Inf | Inf | root                             | free  | angle (rad)                |\n",
    "| 4   | y-orientation of the torso (centre)                                                                             | -Inf | Inf | root                             | free  | angle (rad)                |\n",
    "| 5   | z-orientation of the torso (centre)                                                                             | -Inf | Inf | root                             | free  | angle (rad)                |\n",
    "| 6   | w-orientation of the torso (centre)                                                                             | -Inf | Inf | root                             | free  | angle (rad)                |\n",
    "| 7   | z-angle of the abdomen (in lower_waist)                                                                         | -Inf | Inf | abdomen_z                        | hinge | angle (rad)                |\n",
    "| 8   | y-angle of the abdomen (in lower_waist)                                                                         | -Inf | Inf | abdomen_y                        | hinge | angle (rad)                |\n",
    "| 10  | x-angle of the abdomen (in pelvis)                                                                              | -Inf | Inf | abdomen_x                        | hinge | angle (rad)                |\n",
    "| 11  | x-coordinate of angle between pelvis and right hip (in right_thigh)                                             | -Inf | Inf | right_hip_x                      | hinge | angle (rad)                |\n",
    "| 12  | z-coordinate of angle between pelvis and right hip (in right_thigh)                                             | -Inf | Inf | right_hip_z                      | hinge | angle (rad)                |\n",
    "| 13  | y-coordinate of angle between pelvis and right hip (in right_thigh)                                             | -Inf | Inf | right_hip_y                      | hinge | angle (rad)                |\n",
    "| 14  | angle between right hip and the right shin (in right_knee)                                                      | -Inf | Inf | right_knee                       | hinge | angle (rad)                |\n",
    "| 15  | x-coordinate of angle between pelvis and left hip (in left_thigh)                                               | -Inf | Inf | left_hip_x                       | hinge | angle (rad)                |\n",
    "| 16  | z-coordinate of angle between pelvis and left hip (in left_thigh)                                               | -Inf | Inf | left_hip_z                       | hinge | angle (rad)                |\n",
    "| 17  | y-coordinate of angle between pelvis and left hip (in left_thigh)                                               | -Inf | Inf | left_hip_y                       | hinge | angle (rad)                |\n",
    "| 18  | angle between left hip and the left shin (in left_knee)                                                         | -Inf | Inf | left_knee                        | hinge | angle (rad)                |\n",
    "| 19  | coordinate-1 (multi-axis) angle between torso and right arm (in right_upper_arm)                                | -Inf | Inf | right_shoulder1                  | hinge | angle (rad)                |\n",
    "| 20  | coordinate-2 (multi-axis) angle between torso and right arm (in right_upper_arm)                                | -Inf | Inf | right_shoulder2                  | hinge | angle (rad)                |\n",
    "| 21  | angle between right upper arm and right_lower_arm                                                               | -Inf | Inf | right_elbow                      | hinge | angle (rad)                |\n",
    "| 22  | coordinate-1 (multi-axis) angle between torso and left arm (in left_upper_arm)                                  | -Inf | Inf | left_shoulder1                   | hinge | angle (rad)                |\n",
    "| 23  | coordinate-2 (multi-axis) angle between torso and left arm (in left_upper_arm)                                  | -Inf | Inf | left_shoulder2                   | hinge | angle (rad)                |\n",
    "| 24  | angle between left upper arm and left_lower_arm                                                                 | -Inf | Inf | left_elbow                       | hinge | angle (rad)                |\n",
    "| 25  | x-coordinate velocity of the torso (centre)                                                                     | -Inf | Inf | root                             | free  | velocity (m/s)             |\n",
    "| 26  | y-coordinate velocity of the torso (centre)                                                                     | -Inf | Inf | root                             | free  | velocity (m/s)             |\n",
    "| 27  | z-coordinate velocity of the torso (centre)                                                                     | -Inf | Inf | root                             | free  | velocity (m/s)             |\n",
    "| 28  | x-coordinate angular velocity of the torso (centre)                                                             | -Inf | Inf | root                             | free  | anglular velocity (rad/s)  |\n",
    "| 29  | y-coordinate angular velocity of the torso (centre)                                                             | -Inf | Inf | root                             | free  | anglular velocity (rad/s)  |\n",
    "| 30  | z-coordinate angular velocity of the torso (centre)                                                             | -Inf | Inf | root                             | free  | anglular velocity (rad/s)  |\n",
    "| 31  | z-coordinate of angular velocity of the abdomen (in lower_waist)                                                | -Inf | Inf | abdomen_z                        | hinge | anglular velocity (rad/s)  |\n",
    "| 32  | y-coordinate of angular velocity of the abdomen (in lower_waist)                                                | -Inf | Inf | abdomen_y                        | hinge | anglular velocity (rad/s)  |\n",
    "| 33  | x-coordinate of angular velocity of the abdomen (in pelvis)                                                     | -Inf | Inf | abdomen_x                        | hinge | aanglular velocity (rad/s) |\n",
    "| 34  | x-coordinate of the angular velocity of the angle between pelvis and right hip (in right_thigh)                 | -Inf | Inf | right_hip_x                      | hinge | anglular velocity (rad/s)  |\n",
    "| 35  | z-coordinate of the angular velocity of the angle between pelvis and right hip (in right_thigh)                 | -Inf | Inf | right_hip_z                      | hinge | anglular velocity (rad/s)  |\n",
    "| 36  | y-coordinate of the angular velocity of the angle between pelvis and right hip (in right_thigh)                 | -Inf | Inf | right_hip_y                      | hinge | anglular velocity (rad/s)  |\n",
    "| 37  | angular velocity of the angle between right hip and the right shin (in right_knee)                              | -Inf | Inf | right_knee                       | hinge | anglular velocity (rad/s)  |\n",
    "| 38  | x-coordinate of the angular velocity of the angle between pelvis and left hip (in left_thigh)                   | -Inf | Inf | left_hip_x                       | hinge | anglular velocity (rad/s)  |\n",
    "| 39  | z-coordinate of the angular velocity of the angle between pelvis and left hip (in left_thigh)                   | -Inf | Inf | left_hip_z                       | hinge | anglular velocity (rad/s)  |\n",
    "| 40  | y-coordinate of the angular velocity of the angle between pelvis and left hip (in left_thigh)                   | -Inf | Inf | left_hip_y                       | hinge | anglular velocity (rad/s)  |\n",
    "| 41  | angular velocity of the angle between left hip and the left shin (in left_knee)                                 | -Inf | Inf | left_knee                        | hinge | anglular velocity (rad/s)  |\n",
    "| 42  | coordinate-1 (multi-axis) of the angular velocity of the angle between torso and right arm (in right_upper_arm) | -Inf | Inf | right_shoulder1                  | hinge | anglular velocity (rad/s)  |\n",
    "| 43  | coordinate-2 (multi-axis) of the angular velocity of the angle between torso and right arm (in right_upper_arm) | -Inf | Inf | right_shoulder2                  | hinge | anglular velocity (rad/s)  |\n",
    "| 44  | angular velocity of the angle between right upper arm and right_lower_arm                                       | -Inf | Inf | right_elbow                      | hinge | anglular velocity (rad/s)  |\n",
    "| 45  | coordinate-1 (multi-axis) of the angular velocity of the angle between torso and left arm (in left_upper_arm)   | -Inf | Inf | left_shoulder1                   | hinge | anglular velocity (rad/s)  |\n",
    "| 46  | coordinate-2 (multi-axis) of the angular velocity of the angle between torso and left arm (in left_upper_arm)   | -Inf | Inf | left_shoulder2                   | hinge | anglular velocity (rad/s)  |\n",
    "| 47  | angular velocitty of the angle between left upper arm and left_lower_arm                                        | -Inf | Inf | left_elbow                       | hinge | anglular velocity (rad/s)  |\n",
    "\n",
    "Example code:\n",
    "```\n",
    "step = 0\n",
    "past_rewards = []\n",
    "\n",
    "past_max_torso_position = [-9999]\n",
    "def new_reward(reward):\n",
    "    global step, past_rewards, past_max_torso_position\n",
    "    #x_coor_torso = env.state_vector()[0]\n",
    "    #y_coor_torso = env.state_vector()[1]\n",
    "    z_coor_torso = env.state_vector()[2]\n",
    "    \n",
    "    if z_coor_torso < 0:\n",
    "        salto = z_coor_torso*5\n",
    "    else:\n",
    "        salto = z_coor_torso*20 \n",
    "        \n",
    "    if z_coor_torso > past_max_torso_position[-1]:\n",
    "        past_max_torso_position.append(z_torso)\n",
    "\n",
    "    reward=reward + salto\n",
    "    past_rewards.append(reward) \n",
    "    step += 1\n",
    "    return reward\n",
    "    \n",
    "env = gym.wrappers.TransformReward(env, walk_reward)\n",
    "env.reset()\n",
    "```  \n",
    "\n",
    "EVERY EPISODE IS 1000 STEPS LONG, SO YOU CAN STORE DATA FROM PREVIOUS STEPS FOR THE NEXT STEP\n",
    "Keep the same function prototype which is `new_reward(reward)`. And to get the observations from the table use `env.state_vector()[i]` for the i'th observation of the table.\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "memory = ConversationBufferMemory(\n",
    "    llm=gpt4,\n",
    "    memory_key=\"chat_history\",\n",
    "    return_messages=True,\n",
    "    )\n",
    "            \n",
    "            \n",
    "messages = [\n",
    "SystemMessagePromptTemplate.from_template(system_template),\n",
    "MessagesPlaceholder(variable_name=\"chat_history\"),\n",
    "HumanMessagePromptTemplate.from_template(\"{input}\")\n",
    "]\n",
    "prompt = ChatPromptTemplate.from_messages(messages=messages)\n",
    "\n",
    "chain = LLMChain(\n",
    "llm=gpt4,\n",
    "prompt=prompt,\n",
    "verbose=True,\n",
    "memory=memory,\n",
    ")\n",
    "\n",
    "\n",
    "user_prompt = input(\"> \")\n",
    "#user_prompt = \"Give me a reward function for the humanoid to jump\"\n",
    "output = chain.run({\"input\":user_prompt})\n",
    "# Extract python code (should be wrapped in ```python ```)\n",
    "code = output.split(\"```python\")[1].split(\"```\")[0]\n",
    "exec(code)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rest of your code remains the same\n",
    "if __name__ == '__main__':\n",
    "    gymenv_name = 'Humanoid-v4'\n",
    "    sb3_algo = 'SAC'\n",
    "    path_to_model = './SAC_2500000.zip'\n",
    "\n",
    "    train_model = True\n",
    "    test_model = False\n",
    "\n",
    "    if train_model:\n",
    "        train(env, sb3_algo)\n",
    "\n",
    "    if test_model:\n",
    "        if os.path.isfile(path_to_model):\n",
    "            test(env, sb3_algo, path_to_model=path_to_model)\n",
    "        else:\n",
    "            print(f'{path_to_model} not found.')\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Give feedback to the agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Since it has memory, it remembers the last code and can fix upon it\n",
    "user_prompt = input(\"> \")\n",
    "output = chain.run({\"input\":user_prompt})"
   ]
  }
 ],
 "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.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}