Switch to eqx

examples/env_functions_examples.ipynb

@@ -0,0 +1,202 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e55cd4c2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import jax\n",
+    "import jax.numpy as jnp\n",
+    "from exciting_environments import EnvironmentRegistry\n",
+    "from exciting_environments.utils import MinMaxNormalization\n",
+    "from exciting_environments.pmsm.motor_parameters import MotorVariant"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "76f33ae6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pend_env=EnvironmentRegistry.PENDULUM.make()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "9cc37913",
+   "metadata": {},
+   "source": [
+    "## Step and Simulate ahead"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "18428d21",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "key=jax.random.PRNGKey(1234)\n",
+    "obs, state = pend_env.reset(key)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "0412c1c0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "obs,states,last_state=pend_env.sim_ahead(state,jnp.ones((4,1)))\n",
+    "obs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "68ef9767",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pend_env.generate_rew_trunc_term_ahead(states,jnp.ones((4,1)))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da78e505",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "key=jax.random.PRNGKey(1234)\n",
+    "obs, state = pend_env.reset()\n",
+    "generated_observations = []\n",
+    "generated_actions= []\n",
+    "generated_observations.append(obs)\n",
+    "for i in range(4):\n",
+    "    key,subkey= jax.random.split(key)\n",
+    "    action = jnp.ones(2)#jax.random.uniform(subkey,(2,),minval=-1,maxval=1)\n",
+    "    obs, state = pend_env.step(state, action)\n",
+    "    generated_actions.append(action)\n",
+    "    generated_observations.append(obs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "dc4e3aa7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "generated_observations"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "6c5efe87",
+   "metadata": {},
+   "source": [
+    "### Vmapped"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cdfab95d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "env1=EnvironmentRegistry.PENDULUM.make(static_params={\"g\": jnp.array(9.81), \"l\": jnp.array(1.0),  \"m\": jnp.array(1.0)})\n",
+    "env2=EnvironmentRegistry.PENDULUM.make(static_params={\"g\": jnp.array(9.81), \"l\": jnp.array(2.0),  \"m\": jnp.array(1.0)})\n",
+    "envs = [env1,env2]\n",
+    "batched_envs = EnvironmentRegistry.batch_envs(envs)\n",
+    "batch_size= 2\n",
+    "#batched_envs = EnvironmentRegistry.PENDULUM.make(batch_size=batch_size)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "79b34394",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "keys = jax.random.split(jax.random.PRNGKey(0), batch_size)\n",
+    "obs, states = batched_envs.vmap_reset(keys)\n",
+    "print(obs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c6b42ea8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "state_in=batched_envs.vmap_generate_state_from_observation(obs,keys)\n",
+    "state_in"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b63d4bf8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "actions = jnp.ones((batch_size,1))\n",
+    "obs, states = batched_envs.vmap_reset()\n",
+    "next_obs, next_states = batched_envs.vmap_step(states, actions)\n",
+    "print(next_obs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aa145276",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "keys = jax.random.split(jax.random.PRNGKey(0), batch_size)\n",
+    "obs, states = batched_envs.vmap_reset()\n",
+    "actions = jnp.ones((batch_size,4,1))\n",
+    "next_obs, next_states, last_state = batched_envs.vmap_sim_ahead(states, actions)\n",
+    "print(next_obs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8aa65aea",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "batched_envs.vmap_generate_rew_trunc_term_ahead(next_states, actions)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "venv_dev",
+   "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.11.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

examples/example_gymwrapper.ipynb

@@ -2,52 +2,43 @@
    "cells": [
       {
          "cell_type": "code",
-         "execution_count": 1,
+         "execution_count": null,
          "metadata": {},
          "outputs": [],
          "source": [
-            "import jax.numpy as jnp\n",
             "import jax\n",
-            "import numpy as np\n",
-            "import time\n",
-            "import gymnasium as gym\n",
-            "import sys\n",
-            "sys.path.append(\"..\")\n",
-            "import exciting_environments as excenv\n",
-            "from exciting_environments import GymWrapper\n",
-            "import jax_dataclasses as jdc\n",
-            "from dataclasses import fields\n",
-            "from exciting_environments.utils import MinMaxNormalization\n",
-            "jax.config.update(\"jax_enable_x64\", True)"
+            "import jax.numpy as jnp\n",
+            "from exciting_environments import GymWrapper, EnvironmentRegistry\n",
+            "from exciting_environments.utils import MinMaxNormalization"
          ]
       },
       {
          "cell_type": "code",
-         "execution_count": null,
+         "execution_count": 3,
          "metadata": {},
          "outputs": [],
          "source": [
-            "pend=excenv.PMSM(batch_size=5) # ,action_normalizations={\"torque\":MinMaxNormalization(min=-jnp.array([10,20,30,40,50]),max=jnp.array([10,20,30,40,50]))}\n",
-            "gym_pend=GymWrapper(env=pend,control_state=[]) #,control_state=[\"theta\",\"omega\"]"
+            "pend=EnvironmentRegistry.PENDULUM.make(batch_size=6,control_state=[\"theta\"]) # ,action_normalizations={\"torque\":MinMaxNormalization(min=-jnp.array([10,20,30,40,50]),max=jnp.array([10,20,30,40,50]))}\n",
+            "gym_pend=GymWrapper(env=pend,control_state=[\"theta\"]) #,control_state=[\"theta\",\"omega\"]"
          ]
       },
       {
          "cell_type": "code",
-         "execution_count": 3,
+         "execution_count": 4,
          "metadata": {},
          "outputs": [],
          "source": [
-            "env_key=jax.vmap(jax.random.PRNGKey)(np.random.randint(0, 2**31, size=(pend.batch_size,)))\n",
-            "ref_key=jax.random.PRNGKey(4)"
+            "obs,state=pend.vmap_reset()"
          ]
       },
       {
          "cell_type": "code",
-         "execution_count": null,
+         "execution_count": 5,
          "metadata": {},
          "outputs": [],
          "source": [
-            "obs,_=gym_pend.reset()#rng_env=env_key,rng_ref=ref_key"
+            "env_key=jax.random.split(jax.random.PRNGKey(123),num=6)\n",
+            "ref_key=jax.random.split(jax.random.PRNGKey(223),num=6)"
          ]
       },
       {
@@ -56,27 +47,23 @@
          "metadata": {},
          "outputs": [],
          "source": [
-            "gym_pend.step(action=10*jnp.ones((5,2)))"
+            "gym_pend.reset(rng_ref=ref_key,rng_env=env_key)"
          ]
       },
       {
          "cell_type": "code",
          "execution_count": null,
          "metadata": {},
          "outputs": [],
-         "source": []
-      },
-      {
-         "cell_type": "code",
-         "execution_count": null,
-         "metadata": {},
-         "outputs": [],
-         "source": []
+         "source": [
+            "obs, reward, terminated, truncated = gym_pend.step(action=10*jnp.ones((6,1)))\n",
+            "obs"
+         ]
       }
    ],
    "metadata": {
       "kernelspec": {
-         "display_name": "Python 3",
+         "display_name": "venv_dev",
          "language": "python",
          "name": "python3"
       },
@@ -90,7 +77,7 @@
          "name": "python",
          "nbconvert_exporter": "python",
          "pygments_lexer": "ipython3",
-         "version": "3.11.3"
+         "version": "3.11.9"
       }
    },
    "nbformat": 4,