| --- |
| license: apache-2.0 |
| pipeline_tag: reinforcement-learning |
| tags: |
| - rl |
| - reinforcement-learning |
| - mujoco |
| - physics |
| - balance |
| - mrbalance |
| - custom-architecture |
| --- |
| |
| # MrBalance |
|
|
| Mr. Balance is a plate, and his only goal is to balance whatever object is placed onto him. |
|
|
| <video |
| src="https://huggingface.co/fromziro/MrBalance/resolve/main/assets/video.mp4" |
| controls |
| autoplay |
| loop |
| muted |
| playsinline |
| width="35%"> |
| </video> |
|
|
| Mr. Balance hates when objects fall, so he tries to keep them balanced. |
|
|
| ## Model Architecture |
|
|
| * Base Architecture: `MrBalanceMLPForRL` |
| * Observation Size: `64` |
| * Hidden Size: `128` (first, second and third layer) |
| * Bottleneck Size: `64` (last layer) |
| * Layers: `4` |
| * Action Space: `2` |
|
|
| Mr. Balance uses a shared Actor-Critic MLP designed for continuous-control reinforcement learning. The 64-dimensional observation vector is passed through the four layers with widths of 128 for the first 3 layers and 64 for the last one, using SiLU activations between each of them. The resulting 64-dimensional bottleneck representation is shared by both the actor and critic. The actor produces two control outputs corresponding to the plate's roll and pitch axes, using a Gaussian policy with learned action standard deviations and Tanh squashing to constrain actions to [-1, 1]. The critic uses the same representation to estimate the scalar state value V(s). This shared architecture keeps the policy net small while allowing it to combine object position, velocity, orientation, plate state, and object-property stuff into a combined control representation. |
|
|
| ## Training Configuration |
|
|
| * Number of Environments: `8` |
| * Rollout Steps: `512` |
| * PPO Epochs: `6` |
| * Minibatch Size: `2048` |
| * Gamma: `0.99` |
| * GAE Lambda: `0.95` |
| * Clip Coefficient: `0.20` |
| * Value Coefficient: `0.50` |
| * Entropy Coefficient: `0.005` |
| * Maximum Gradient Norm: `0.50` |
| * Learning Rate: `3e-4` |
| * Adam Epsilon: `1e-5` |
| * Anneal Learning Rate: `true` |
| * Number of Episodes: `12756` |
| * Maximum Episode Steps: `10000` |
|
|
| ### Training Objects |
|
|
| - Sphere |
| - Egg |
| - Heavy Ball |
|
|
| ## Training Results |
|
|
| | Object | Reward Mean | Len Mean | Survival % | Tracking Err | |
| | ---------- | ----------: | -------: | ---------: | -----------: | |
| | sphere | 24,398.48 | 10,000.0 | 100.0% | 0.0205m | |
| | egg | 23,240.42 | 10,000.0 | 100.0% | 0.0458m | |
| | heavy_ball | 24,310.68 | 10,000.0 | 100.0% | 0.0401m | |
| |
| Mr. Balance excels at balancing all the objects he was trained on, but does he generalize to other objects? |
| |
| | Object | Reward Mean | Len Mean | Survival % | Tracking Err | |
| | --------------- | ----------: | -------: | ---------: | -----------: | |
| | sphere | 24,398.48 | 10,000.0 | 100.0% | 0.0205m | |
| | disk | 21,191.56 | 10,000.0 | 100.0% | 0.1362m | |
| | egg | 23,240.42 | 10,000.0 | 100.0% | 0.0458m | |
| | cup | 21,702.17 | 10,000.0 | 100.0% | 0.1242m | |
| | coin | 20,973.76 | 10,000.0 | 100.0% | 0.1428m | |
| | stick | 13,601.42 | 8,477.4 | 80.0% | 0.3399m | |
| | tall | 22,265.77 | 10,000.0 | 100.0% | 0.1077m | |
| | triangle | 18,794.99 | 10,000.0 | 100.0% | 0.2181m | |
| | block | 21,489.38 | 10,000.0 | 100.0% | 0.1316m | |
| | puck | 21,265.33 | 10,000.0 | 100.0% | 0.1361m | |
| | cone | 18,327.38 | 10,000.0 | 100.0% | 0.2316m | |
| | capsule | 22,302.92 | 10,000.0 | 100.0% | 0.1097m | |
| | wedge | 18,921.33 | 10,000.0 | 100.0% | 0.2148m | |
| | tetra | 18,838.96 | 10,000.0 | 100.0% | 0.2165m | |
| | flat_bar | 20,550.93 | 10,000.0 | 100.0% | 0.1526m | |
| | cross | 21,394.63 | 10,000.0 | 100.0% | 0.1435m | |
| | L_shape | 21,221.41 | 10,000.0 | 100.0% | 0.1389m | |
| | wide_block | 20,857.84 | 10,000.0 | 100.0% | 0.1611m | |
| | heavy_ball | 24,310.68 | 10,000.0 | 100.0% | 0.0401m | |
| | offcenter_block | 21,408.22 | 10,000.0 | 100.0% | 0.1395m | |
|
|
| The answer is yes; Mr. Balance excels at balancing a wide range of objects, not just the three he was trained on. |
|
|
| ## Inference |
|
|
| First install the required dependencies: |
|
|
| ```bash |
| pip install torch transformers safetensors "mujoco==3.10.0" numpy |
| ``` |
|
|
| Windows users also might need the Microsoft C++ Build Tools: |
|
|
| [https://visualstudio.microsoft.com/visual-cpp-build-tools/](https://visualstudio.microsoft.com/visual-cpp-build-tools/) |
|
|
| Next, download both `inference.py` and `balance_plate_rl.py` and run: |
|
|
| ```bash |
| python inference.py |
| ``` |
|
|
| Add `--render` for live visual rendering or/and `--object` to choose a specific object. |
|
|
| ## License |
|
|
| Apache 2.0. |
|
|
| ## Citation |
|
|
| ``` |
| |
| ``` |
