update notebook (#39)

Author: WT-MM

train.ipynb

@@ -43,7 +43,7 @@
    "source": [
     "# Install packages\n",
     "\n",
-    "!pip install ksim==0.1.0 xax==0.3.0 mujoco-scenes"
+    "!pip install ksim==0.1.2 xax==0.3.0 mujoco-scenes"
    ]
   },
   {
@@ -474,10 +474,6 @@
     "        value=True,\n",
     "        help=\"Whether to use the IMU acceleration and gyroscope observations.\",\n",
     "    )\n",
-    "    use_domain_randomization: bool = xax.field(\n",
-    "        value=True,\n",
-    "        help=\"Whether to use domain randomization.\",\n",
-    "    )\n",
     "\n",
     "    # Curriculum parameters.\n",
     "    num_curriculum_levels: int = xax.field(\n",
@@ -628,22 +624,20 @@
     "    def get_rewards(self, physics_model: ksim.PhysicsModel) -> list[ksim.Reward]:\n",
     "        return [\n",
     "            # Standard rewards.\n",
-    "            ksim.NaiveForwardReward(clip_max=2.0, in_robot_frame=False, scale=3.0),\n",
+    "            ksim.NaiveForwardReward(clip_max=1.25, in_robot_frame=False, scale=3.0),\n",
     "            ksim.NaiveForwardOrientationReward(scale=1.0),\n",
     "            ksim.StayAliveReward(scale=1.0),\n",
     "            ksim.UprightReward(scale=0.5),\n",
     "            # Avoid movement penalties.\n",
-    "            ksim.AngularVelocityPenalty(index=(\"x\", \"y\"), scale=-0.005),\n",
-    "            ksim.LinearVelocityPenalty(index=(\"z\"), scale=-0.005),\n",
+    "            ksim.AngularVelocityPenalty(index=(\"x\", \"y\"), scale=-0.1),\n",
+    "            ksim.LinearVelocityPenalty(index=(\"z\"), scale=-0.1),\n",
     "            # Normalization penalties.\n",
     "            ksim.AvoidLimitsPenalty.create(physics_model, scale=-0.01),\n",
-    "            ksim.JointVelocityPenalty(scale=-0.01, scale_by_curriculum=True),\n",
     "            ksim.JointAccelerationPenalty(scale=-0.01, scale_by_curriculum=True),\n",
     "            ksim.JointJerkPenalty(scale=-0.01, scale_by_curriculum=True),\n",
     "            ksim.LinkAccelerationPenalty(scale=-0.01, scale_by_curriculum=True),\n",
     "            ksim.LinkJerkPenalty(scale=-0.01, scale_by_curriculum=True),\n",
     "            ksim.ActionAccelerationPenalty(scale=-0.01, scale_by_curriculum=True),\n",
-    "            ksim.CtrlPenalty(scale=-0.01, scale_by_curriculum=True),\n",
     "            # Bespoke rewards.\n",
     "            BentArmPenalty.create_penalty(physics_model, scale=-0.1),\n",
     "            StraightLegPenalty.create_penalty(physics_model, scale=-0.1),\n",
@@ -656,11 +650,8 @@
     "        ]\n",
     "\n",
     "    def get_curriculum(self, physics_model: ksim.PhysicsModel) -> ksim.Curriculum:\n",
-    "        return ksim.ConstantCurriculum(\n",
-    "            # We toggle domain randomization by setting the curriculum level.\n",
-    "            # Since the domain randomization functions all use this level,\n",
-    "            # this effectively toggles them on and off.\n",
-    "            level=1.0 if self.config.use_domain_randomization else 0.0,\n",
+    "        return ksim.DistanceFromOriginCurriculum(\n",
+    "            min_level_steps=5,\n",
     "        )\n",
     "\n",
     "    def get_model(self, key: PRNGKeyArray) -> Model:\n",