Merge remote-tracking branch 'upstream/main' into sensor_touch

and updates to utilize ray geom functionality with touch sensor

Author: thowell

.github/workflows/ci.yml

@@ -70,3 +70,22 @@ jobs:
       - name: Test with pytest
         run: |
           pytest
+
+  kernel_analyzer:
+    name: Kernel analyzer
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install mujoco --pre -f https://py.mujoco.org/
+          pip install warp-lang --pre --upgrade -f https://pypi.nvidia.com/warp-lang/
+          pip install -e .[dev,cpu]
+      - name: Run kernel analyzer
+        run: |
+          python contrib/kernel_analyzer/kernel_analyzer/cli.py mujoco_warp/_src/*.py --output=github

.gitignore

@@ -208,4 +208,7 @@ Temporary Items
 .apdisk
 
 # Node stuff
-node_modules
+node_modules
+
+# VS Code
+.vscode/

.pre-commit-config.yaml

@@ -0,0 +1,18 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.4.2 # Use the latest Ruff version
+    hooks:
+      - id: ruff
+        args: [--fix, --exit-non-zero-on-fix]
+      - id: ruff-format
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.6.0 # Use the latest version
+    hooks:
+      - id: check-yaml
+  - repo: local
+    hooks:
+      - id: kernel-analyzer
+        name: kernel-analyzer
+        entry: python contrib/kernel_analyzer/kernel_analyzer/cli.py
+        language: system
+        types: [python]

README.md

@@ -39,6 +39,10 @@ pytest
 
 Should print out something like `XX passed in XX.XXs` at the end!
 
+If you plan to write Warp kernels for MJWarp, please use the `kernel_analyzer` vscode plugin located in `contrib/kernel_analyzer`.
+Please see the `README.md` there for details on how to install it and use it.  The same kernel analyzer will be run on any PR
+you open, so it's important to fix any issues it reports.
+
 # Compatibility
 
 The following features are implemented:

conftest.py

@@ -17,9 +17,7 @@
 
 
 def pytest_addoption(parser):
-  parser.addoption(
-    "--cpu", action="store_true", default=False, help="run tests with cpu"
-  )
+  parser.addoption("--cpu", action="store_true", default=False, help="run tests with cpu")
   parser.addoption(
     "--verify_cuda",
     action="store_true",

contrib/apptronik_apollo_locomotion.ipynb

@@ -176,8 +176,8 @@
     "  params[\"site_xpos\"] = True\n",
     "  out_batched = mjx.Data(**params)\n",
     "\n",
-    "  qpos, qvel, xpos, xmat, qacc_warmstart, subtree_com, cvel, site_xpos = (\n",
-    "    jax_mjwarp_step(d.ctrl, d.qpos, d.qvel, d.qacc_warmstart)\n",
+    "  qpos, qvel, xpos, xmat, qacc_warmstart, subtree_com, cvel, site_xpos = jax_mjwarp_step(\n",
+    "    d.ctrl, d.qpos, d.qvel, d.qacc_warmstart\n",
     "  )\n",
     "  d = d.replace(\n",
     "    qpos=qpos,\n",
@@ -355,9 +355,7 @@
     "    obs = self._get_obs(data, state.info, contact)\n",
     "    done = self._get_termination(data)\n",
     "\n",
-    "    rewards = self._get_reward(\n",
-    "      data, action, state.info, state.metrics, done, first_contact, contact\n",
-    "    )\n",
+    "    rewards = self._get_reward(data, action, state.info, state.metrics, done, first_contact, contact)\n",
     "    rewards = {k: v * self._config.reward_config.scales[k] for k, v in rewards.items()}\n",
     "    reward = sum(rewards.values()) * self.dt\n",
     "\n",
@@ -390,9 +388,7 @@
     "    fall_termination = data.xpos[self._head_body_id, 2] < 1.0\n",
     "    return fall_termination\n",
     "\n",
-    "  def _get_obs(\n",
-    "    self, data: mjx.Data, info: dict[str, Any], contact: jax.Array\n",
-    "  ) -> mjx_env.Observation:\n",
+    "  def _get_obs(self, data: mjx.Data, info: dict[str, Any], contact: jax.Array) -> mjx_env.Observation:\n",
     "    cos = jp.cos(info[\"phase\"])\n",
     "    sin = jp.sin(info[\"phase\"])\n",
     "    phase = jp.concatenate([cos, sin])\n",
@@ -424,36 +420,24 @@
     "    del metrics  # Unused.\n",
     "    return {\n",
     "      # Tracking rewards.\n",
-    "      \"tracking_lin_vel\": self._reward_tracking_lin_vel(\n",
-    "        info[\"command\"], self._get_global_linvel(data, self._torso_id)\n",
-    "      ),\n",
-    "      \"tracking_ang_vel\": self._reward_tracking_ang_vel(\n",
-    "        info[\"command\"], self._get_global_angvel(data, self._torso_id)\n",
-    "      ),\n",
+    "      \"tracking_lin_vel\": self._reward_tracking_lin_vel(info[\"command\"], self._get_global_linvel(data, self._torso_id)),\n",
+    "      \"tracking_ang_vel\": self._reward_tracking_ang_vel(info[\"command\"], self._get_global_angvel(data, self._torso_id)),\n",
     "      # Base-related rewards.\n",
-    "      \"ang_vel_xy\": self._cost_ang_vel_xy(\n",
-    "        self._get_global_angvel(data, self._torso_id)\n",
-    "      ),\n",
+    "      \"ang_vel_xy\": self._cost_ang_vel_xy(self._get_global_angvel(data, self._torso_id)),\n",
     "      \"orientation\": self._cost_orientation(self._get_z_frame(data, self._torso_id)),\n",
     "      # Energy related rewards.\n",
-    "      \"action_rate\": self._cost_action_rate(\n",
-    "        action, info[\"last_act\"], info[\"last_last_act\"]\n",
-    "      ),\n",
+    "      \"action_rate\": self._cost_action_rate(action, info[\"last_act\"], info[\"last_last_act\"]),\n",
     "      # Feet related rewards.\n",
     "      \"feet_slip\": self._cost_feet_slip(data, contact, info),\n",
-    "      \"feet_air_time\": self._reward_feet_air_time(\n",
-    "        info[\"feet_air_time\"], first_contact, info[\"command\"]\n",
-    "      ),\n",
+    "      \"feet_air_time\": self._reward_feet_air_time(info[\"feet_air_time\"], first_contact, info[\"command\"]),\n",
     "      \"feet_phase\": self._reward_feet_phase(\n",
     "        data,\n",
     "        info[\"phase\"],\n",
     "        self._config.reward_config.max_foot_height,\n",
     "        info[\"command\"],\n",
     "      ),\n",
     "      # Pose related rewards.\n",
-    "      \"joint_deviation_hip\": self._cost_joint_deviation_hip(\n",
-    "        data.qpos[7:], info[\"command\"]\n",
-    "      ),\n",
+    "      \"joint_deviation_hip\": self._cost_joint_deviation_hip(data.qpos[7:], info[\"command\"]),\n",
     "      \"joint_deviation_knee\": self._cost_joint_deviation_knee(data.qpos[7:]),\n",
     "      \"dof_pos_limits\": self._cost_joint_pos_limits(data.qpos[7:]),\n",
     "      \"pose\": self._cost_pose(data.qpos[7:]),\n",
@@ -504,17 +488,13 @@
     "\n",
     "  # Energy related rewards.\n",
     "\n",
-    "  def _cost_action_rate(\n",
-    "    self, act: jax.Array, last_act: jax.Array, last_last_act: jax.Array\n",
-    "  ) -> jax.Array:\n",
+    "  def _cost_action_rate(self, act: jax.Array, last_act: jax.Array, last_last_act: jax.Array) -> jax.Array:\n",
     "    del last_last_act  # Unused.\n",
     "    return jp.sum(jp.square(act - last_act))\n",
     "\n",
     "  # Feet related rewards.\n",
     "\n",
-    "  def _cost_feet_slip(\n",
-    "    self, data: mjx.Data, contact: jax.Array, info: dict[str, Any]\n",
-    "  ) -> jax.Array:\n",
+    "  def _cost_feet_slip(self, data: mjx.Data, contact: jax.Array, info: dict[str, Any]) -> jax.Array:\n",
     "    del info  # Unused.\n",
     "    body_vel = self._get_global_linvel(data, self._torso_id)[:2]\n",
     "    reward = jp.sum(jp.linalg.norm(body_vel, axis=-1) * contact)\n",
@@ -534,9 +514,7 @@
     "    reward = jp.sum(air_time)\n",
     "    return reward\n",
     "\n",
-    "  def get_rz(\n",
-    "    phi: Union[jax.Array, float], swing_height: Union[jax.Array, float] = 0.08\n",
-    "  ) -> jax.Array:\n",
+    "  def get_rz(phi: Union[jax.Array, float], swing_height: Union[jax.Array, float] = 0.08) -> jax.Array:\n",
     "    def cubic_bezier_interpolation(y_start, y_end, x):\n",
     "      y_diff = y_end - y_start\n",
     "      bezier = x**3 + 3 * (x**2 * (1 - x))\n",
@@ -556,9 +534,7 @@
     "  ) -> jax.Array:\n",
     "    # Reward for tracking the desired foot height.\n",
     "    foot_pos = data.site_xpos[self._feet_site_id]\n",
-    "    foot_pos = jp.array(\n",
-    "      [jp.mean(foot_pos[0:4], axis=0), jp.mean(foot_pos[4:8], axis=0)]\n",
-    "    )\n",
+    "    foot_pos = jp.array([jp.mean(foot_pos[0:4], axis=0), jp.mean(foot_pos[4:8], axis=0)])\n",
     "    foot_z = foot_pos[..., -1]\n",
     "    rz = Joystick.get_rz(phase, swing_height=foot_height)\n",
     "    error = jp.sum(jp.square(foot_z - rz))\n",
@@ -606,12 +582,8 @@
     "  def sample_command(self, rng: jax.Array) -> jax.Array:\n",
     "    rng1, rng2, rng3, rng4 = jax.random.split(rng, 4)\n",
     "\n",
-    "    lin_vel_x = jax.random.uniform(\n",
-    "      rng1, minval=self._config.lin_vel_x[0], maxval=self._config.lin_vel_x[1]\n",
-    "    )\n",
-    "    lin_vel_y = jax.random.uniform(\n",
-    "      rng2, minval=self._config.lin_vel_y[0], maxval=self._config.lin_vel_y[1]\n",
-    "    )\n",
+    "    lin_vel_x = jax.random.uniform(rng1, minval=self._config.lin_vel_x[0], maxval=self._config.lin_vel_x[1])\n",
+    "    lin_vel_y = jax.random.uniform(rng2, minval=self._config.lin_vel_y[0], maxval=self._config.lin_vel_y[1])\n",
     "    ang_vel_yaw = jax.random.uniform(\n",
     "      rng3,\n",
     "      minval=self._config.ang_vel_yaw[0],\n",
@@ -704,9 +676,7 @@
     "network_factory = ppo_networks.make_ppo_networks\n",
     "if \"network_factory\" in ppo_params:\n",
     "  del ppo_training_params[\"network_factory\"]\n",
-    "  network_factory = functools.partial(\n",
-    "    ppo_networks.make_ppo_networks, **ppo_params.network_factory\n",
-    "  )\n",
+    "  network_factory = functools.partial(ppo_networks.make_ppo_networks, **ppo_params.network_factory)\n",
     "\n",
     "train_fn = functools.partial(\n",
     "  ppo.train,\n",

contrib/kernel_analyzer/README.md

@@ -5,13 +5,13 @@ Kernel Analyzer checks warp kernels to ensure correctness and conformity.  It co
 # CLI usage
 
 ```bash
-python contrib/kernel_analyzer/kernel_analyzer/cli.py --files somefile.py --types mujoco_warp/_src/types.py 
+python contrib/kernel_analyzer/kernel_analyzer/cli.py mujoco_warp/_src/*.py --types mujoco_warp/_src/types.py 
 ```
 
 # CLI for github CI
 
 ```bash
-python contrib/kernel_analyzer/kernel_analyzer/cli.py --files somefile.py --types mujoco_warp/_src/types.py 
+python contrib/kernel_analyzer/kernel_analyzer/cli.py mujoco_warp/_src/*.py --types mujoco_warp/_src/types.py --output=github
 ```
 
 # VSCode plugin
@@ -21,16 +21,17 @@ Enjoy kernel analysis directly within vscode.
 ## Installing kernel analyzer
 
 1. Create a new python env (`python3 -m venv env`) or use your existing mjwarp env (`source env/bin/activate`).
-2. Within the python env, install the kernel analyzer python dependencies:
+2. Within the python env, install the kernel analyzer's python dependencies by pip installing MJWarp dev:
   ```bash
-  pip install -r contrib/kernel_analyzer/kernel_analyzer/requirements.txt
+  cd mujoco_warp
+  pip install -e .[dev]
   ```
 3. Inside vscode, navigate to `contrib/kernel_analyzer/`
 4. Right click on `kernel-analyzer-{version}.vsix` file
 5. Select "Install Extension VSIX"
 6. Open vscode settings and navigate to `Extensions > Kernel Analyzer`
-7. Set **Python Path** to the `bin/python` of the env you set up in step 1, e.g. `/home/$USER/work/mujoco_warp/env/bin/python`
-8. Set **Types Path** to the location of `types.py` in your checked out code, e.g. `/home/$USER/work/mujoco_warp/mujoco_warp/_src/types.py`
+7. Set Python Path to the `bin/python` of the env you set up in step 1, e.g. `/home/$USER/work/mujoco_warp/env/bin/python`
+8. Set Types Path to the location of `types.py` in your checked out code, e.g. `/home/$USER/work/mujoco_warp/mujoco_warp/_src/types.py`
 
 ## Plugin Development
 
@@ -60,4 +61,4 @@ Create a debug configuration in `.vscode/launch.json`:
 
 ```bash
 npm run package
-```
+```

contrib/kernel_analyzer/client/package.json

@@ -1,7 +1,7 @@
 {
   "name": "kernel-analyzer-client",
   "description": "Client-side code (TypeScript) for the Kernel Analyzer VS Code extension",
-  "version": "0.1.0",
+  "version": "0.2.0",
   "private": true,
   "license": "Apache-2.0",
   "engines": {