Wip: Box box

.gitignore

@@ -2,6 +2,7 @@
 *.egg-info/
 env/
 .history
+.vscode
 
 # Python byte-compiled / optimized / DLL files
 __pycache__/

AUTHORS

@@ -0,0 +1,10 @@
+# This is the list of significant contributors to mjWarp.
+#
+# This does not necessarily list everyone who has contributed code,
+# especially since many employees of one corporation may be contributing.
+# To see the full list of contributors, see the revision history in
+# source control.
+
+Google LLC
+NVIDIA Corporation
+

CONTRIBUTING.md

@@ -0,0 +1,25 @@
+# How to Contribute
+
+## Contributor License Agreement
+
+Contributions to this project must be accompanied by a Contributor License
+Agreement. You (or your employer) retain the copyright to your contribution,
+this simply gives us permission to use and redistribute your contributions as
+part of the project. Head over to <https://cla.developers.google.com/> to see
+your current agreements on file or to sign a new one.
+
+You generally only need to submit a CLA once, so if you've already submitted one
+(even if it was for a different project), you probably don't need to do it
+again.
+
+## Code reviews
+
+All submissions, including submissions by project members, require review. We
+use GitHub pull requests for this purpose. Consult
+[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more
+information on using pull requests.
+
+## Community Guidelines
+
+This project follows [Google's Open Source Community
+Guidelines](https://opensource.google/conduct/).

README.md

@@ -1,6 +1,6 @@
-# mjWarp
+# MJWarp
 
-MuJoCo implemented in Warp.
+MuJoCo implemented in NVIDIA Warp.
 
 # Installing for development
 
@@ -32,47 +32,62 @@ Should print out something like `XX passed in XX.XXs` at the end!
 Benchmark as follows:
 
 ```bash
-mjx-testspeed --function=forward --is_sparse=True --mjcf=humanoid/humanoid.xml --batch_size=8192
+mjwarp-testspeed --function=step --mjcf=humanoid/humanoid.xml --batch_size=8192
 ```
 
-Some relevant benchmarks on an NVIDIA GeForce RTX 4090:
+To get a full trace of the physics steps (e.g. timings of the subcomponents) run the following:
 
-## forward steps / sec (smooth dynamics only)
-
-27 dofs per humanoid, 8k batch size.
-
-| Num Humanoids   | MJX    |  mjWarp dense |  mjWarp sparse |
-| ----------------| -------| ------------- | -------------- |
-| 1               | 7.9M   |  15.6M        | 13.7M          |
-| 2               | 2.6M   |  7.4M         | 7.8M           |
-| 3               | 2.2M   |  4.6M         | 5.3M           |
-| 4               | 1.5M   |  3.3M         | 4.1M           |
-| 5               | 1.1M   |  ❌           | 3.2M           |
-
-# Ideas for what to try next
-
-## 1. Unroll steps
-
-In the Pure JAX benchmark, we can tell JAX to unroll some number of FK steps (in the benchmarks above, `unroll=4`).  This has a big impact on performance.  If we change `unroll` from 4 to 1, pure JAX performance at 8k batch drops from 50M to 33M steps/sec.
-
-Is there some way that we can improve Warp performance in the same way?  If I know ahead of time that I am going to call FK in a loop 1000 times, can I somehow inject unroll primitives?
-
-## 2. Different levels of parallelism
-
-The current approach parallelizes over body kinematic tree depth.  We could go either direction: remove body parallism (fewer kernel launches), or parallelize over joints instead (more launches, more parallelism).
-
-## 3. Tiling
+```bash
+mjwarp-testspeed --function=step --mjcf=humanoid/humanoid.xml --batch_size=8192 --event_trace=True
+```
 
-It looks like a thing!  Should we use it?  Will it help?
+`humanoid.xml` has been carefully optimized for MJX in the following ways:
 
-## 4. Quaternions
+* Newton solver iterations are capped at 1, linesearch iterations capped at 4
+* Only foot<>floor collisions are turned on, producing at most 8 contact points
+* Adding a damping term in the Euler integrator (which invokes another `factor_m` and `solve_m`) is disabled
 
-Why oh why did Warp make quaternions x,y,z,w?  In order to be obstinate I wrote my own quaternion math.  Is this slower than using the Warp quaternion primitives?
+By comparing MJWarp to MJX on this model, we are comparing MJWarp to the very best that MJX can do.
 
-## 5. `wp.static`
+For many (most) MuJoCo models, particularly ones that haven't been carefully tuned, MJX will
+do much worse.
 
-Haven't tried this at all - curious to see if it helps.
+## physics steps / sec
 
-## 6. Other stuff?
+NVIDIA GeForce RTX 4090, 27 dofs, ncon=8, 8k batch size.
 
-Should I be playing with `block_dim`?  Is my method for timing OK or did I misunderstand how `wp.synchronize` works?  Is there something about allocating that I should be aware of?  What am I not thinking of?
+```
+Summary for 8192 parallel rollouts
+
+ Total JIT time: 0.82 s
+ Total simulation time: 2.98 s
+ Total steps per second: 2,753,173
+ Total realtime factor: 13,765.87 x
+ Total time per step: 363.22 ns
+
+Event trace:
+
+step: 361.41                 (MJX: 316.58 ns)
+  forward: 359.15
+    fwd_position: 52.58
+      kinematics: 19.36      (MJX:  16.45 ns)
+      com_pos: 7.80          (MJX:  12.37 ns)
+      crb: 12.44             (MJX:  27.91 ns)
+      factor_m: 6.40         (MJX:  27.48 ns)
+      collision: 4.07        (MJX:   1.23 ns)
+      make_constraint: 6.32  (MJX:  42.39 ns)
+      transmission: 1.30     (MJX:   3.54 ns)
+    fwd_velocity: 26.52
+      com_vel: 8.44          (MJX:   9.38 ns)
+      passive: 1.06          (MJX:   3.22 ns)
+      rne: 10.96             (MJX:  16.75 ns)
+    fwd_actuation: 2.74      (MJX:   3.93 ns)
+    fwd_acceleration: 11.90
+      xfrc_accumulate: 3.83  (MJX:   6.81 ns)
+      solve_m: 6.92          (MJX:   8.88 ns)
+    solve: 264.38            (MJX: 153.29 ns)
+      mul_m: 5.93
+      _linesearch_iterative: 43.15
+        mul_m: 3.66
+  euler: 1.74               (MJX:    3.78 ns)
+```

contrib/README.md

@@ -0,0 +1,5 @@
+# Contrib
+
+Contrib is a home for experiments, helper scripts and other tools that are not officially part of MJWarp.
+
+The contents of this directory are subject to change with no notice.