← Back to Write-ups

Mamba + LNN Robot • Part II

Hardening the Robot Brain

By David Nash · · 13 min read (2,752 words)
Total Sparks: 0
Enjoyed this write-up? Charge it with sparks

In Part I: Building a Robot Brain, I trained a 7-DOF Franka Panda arm using Mamba State Space Models and Liquid Neural Networks, reaching a 98.7% success rate at picking up colored cubes and sorting them onto target pads.

That 98.7% headline number looked great, but it only held on the one table I trained on. As soon as I swapped in a different table texture, it fell apart.

This failure is because reinforcement learning algorithms do not learn general robotics principles on their own. They optimize ruthlessly against the exact world you put in front of them. In my initial setup, that world was a mostly static and highly idealized simulation.

When a policy is highly fragile like this, the root causes usually fall into two categories:

  • Blatant Exclusions (Simulation Cheat Codes): Assuming motors respond with zero latency, have infinite acceleration, or the environment will always have perfect lighting. The neural network will optimize with these unrealistic behaviors since it has no way to know any better. It can only optimize based on what it is given.
  • Accidental Exclusions (Unmodeled Blind Spots): Forgetting to model rubber fingertip friction, ignoring joint damping, or not accounting for camera sensor noise. Missing these can cause the policy to learn unnatural compensatory tricks or have unrealistic expectations that break when deployed in reality.

Fixing this means tackling both sides of the problem. I had to model the physical constraints I knew about, and randomize the visual variations I could not predict.

Fixing the Contact Physics

In my Part I setup, the gripper fingertips used 's default contact settings. Under default settings, bodies interact as rigid surfaces with basic sliding friction and zero resistance to twisting.

This was an accidental exclusion that distorted the robot's learning. In the real world, a Franka Panda has soft silicone pads on its fingers. When the fingers close on a cube, the rubber compresses and conforms to the surface, generating additional friction that stops the object from twisting and keeps it from sliding out of the gripper.

Because the default simulation physics gave the gripper zero torsional friction, any off-center grasp made picking up the cube similar to trying to pick up a marble with chopsticks: the slightest misalignment and the thing twists and falls out of the grip. This made the robot learn some less than ideal gripping behavior and the performance of the held grips was overall quite poor. The problem was not the neural network on this one, this one was purely a modeling problem.

To restore realistic contact physics, I updated four key contact properties on the fingertip geoms. The values are silicone-like numbers I chose and tuned, not measurements of the real Franka pads:

Contact PropertyDefault SimulationRubber Pad (tuned)What It Changes in Practice
Contact Friction ModeSliding only 3 dimensionsSliding + Torsional 4 dimensionsEnables torsional friction around the contact normal, stopping off-center cubes from spinning out of the fingers.
Friction Coefficients
Slide: 1.0
Twist: 0.005
Roll: 0.0001
Slide: 1.8 (+80%)
Twist: 0.08 (16×)
Roll: 0.001 (10×)
Increases sliding grip and vastly boosts torsional and rolling friction to match physical silicone rubber.
Contact Softness & DampingSofter / slower response 20 ms, critical damping (1.0)Stiffer, fast response 8 ms, critical damping (1.0)Shortens the contact time constant, giving a stiffer but still critically damped spring that absorbs finger closure impacts without bouncing.
Material Override PriorityAveraged (0) Finger & cube properties are mixedOverride (1) Fingertip rubber dominates contactForces the fingertip's rubber properties to override the cube's defaults instead of averaging them together.

Beyond static friction, physical rubber pads exhibit progressive stiffening. They deform easily on initial touch to conform around the object, then rapidly stiffen as they compress under load. In MuJoCo, this is modeled with a mathematical function called solver impedance (abbreviated in the simulator as solimp).

Rather than treating contact as an all-or-nothing rigid collision, the solimp function dynamically calculates contact stiffness as a function of penetration depth: solimp="dmin dmax width midpoint power". Tuning this to "0.9 0.99 0.003 0.5 2" ramps stiffness up over about 3 mm of penetration, which is roughly the squish I'd expect from a rubber pad. The function starts compliant (0.9) to absorb initial closure impacts without bouncing, then smoothly ramps stiffness up to near-rigid (0.99) under full squeeze so the cube stays locked in place without sinking through the finger mesh.

Fixing this exclusion gave our simulation robot a big upgrade. With proper rubber friction, the cube stopped slipping, allowing the policy to develop smooth, compliant grasping habits. I have not tried it on real hardware.

Scaling Training with Vectorized GPU Physics (MJX)

Training a RL model is a pretty computationally heavy task, even for a lightweight model like the one I've built. Hardening that policy comes with additional computational cost. When you add compliant contacts, motor delays, and wider physical variations, the policy needs tens of millions of environment steps to explore and converge.

On a laptop CPU, running vision-based reinforcement learning was painfully slow. Stepping the physics engine, generating camera frames, and transferring tensors back and forth between CPU memory and GPU memory meant that my training runs were taking around 4-8 hours, just to test a tiny incremental change. The worst part would be that I made a mistake somewhere along the way and that whole training run would just become throwaway work. I quickly realized that it would take too long to learn and adapt trying to do it the free way on my laptop.

To make large-scale hardening practical, I brought out my wallet and moved simulation entirely onto the GPU using .

Rather than running one simulation at a time on the CPU, MJX compiles the entire physics pipeline into fused GPU kernels with JAX. Paired with a hardware-accelerated GPU raytracer, MJX steps physics and renders cameras for 256 parallel environments simultaneously directly in GPU VRAM on an RTX 3090 whereas on my laptop I was only able to run 4 parallel envs at a time.

Because the simulation states, sensor observations, neural network weights, and PPO gradients all live on the same GPU, there is zero memory transfer bottleneck. A training run that once took 8 hours on a laptop to gather 9 million steps completes about 15 million steps, rendering included, in under 45 minutes on a single cloud GPU. That's roughly 18 times the throughput, for under a dollar per hour of compute.

Visual Robustness & Domain Randomization

Now that the physics and training throughput are fixed, the next point of failure is vision.

In Part I, my policy was trained on a single, uniform tabletop surface. I had originally messed up the texture of the table and the wood grain look was warping horribly around the edges of the table but I didn't care too much at first, I just wanted something that worked. The vision encoder learned to spot the colored cube in the scene with this broken texture and its movements relied on it.

Once it was working pretty well I wanted to record a better demo with a table that actually, ya know, looked like a table. So, I swapped the plain rectangular prism table with a wonky texture for a fixed woodgrain texture and gave the table legs. I checked the scene locally without running the sim just to make sure it all looked good, and it did. I was excited to record my new fancy demo.

Surprised Pikachu meme

This was my face when I watched the new recorded demo video because my robot which was previously working well dropped from 98.7% success rate down to 5.7%. The recording was terrible!

Fig. 1 Visual Overfitting Under a Texture Swap

The baseline policy successfully identifying the cube, descending, and placing it when operating on the single table surface it saw during training.

Testing the un-hardened Part I policy against a new woodgrain table texture without retraining. Left: the policy succeeds with the table it was trained on. Right: the exact same policy fails when a new table is swapped in.

This failure illustrates a fundamental rule of deep learning in robotics: neural networks will always exploit anything you let them exploit.

The vision encoder had not generically learned the spatial concept of "a cube resting on a surface." Instead, it had learned to exploit its visual data to identify the specific pixel arrangement of the scene it was trained on and its weights were set based off of that. It didn't need to work any harder to generalize, because I hadn't given it a reason to. When the new table altered the look of the pixels around the cube, the spatial feature maps became noise, and the policy freaked out.

Native 3D Domain Randomization

The standard remedy for visual brittleness is . If the policy sees hundreds of visual variations during training, it learns to ignore irrelevant background features and instead focus on what actually matters: physical object edges, depth cues, and relative spatial geometry.

Doing domain randomization correctly requires some forethought, though. You need to add randomization in a way that makes sense for your scene. Think about what the robot might actually encounter in its real world operating environment. For a tabletop arm, that means lighting and surface textures. You need to randomize those so it becomes robust against scenarios not explicitly covered in its training data. I only randomized what I expected the arm to see, since randomizing everything indiscriminately can drown out the learning signal.

The key for me was to randomize a few things natively in the 3D scene. I assigned the table in each of the 256 parallel training environments one of eight real-world surface materials—including wood, concrete, linoleum, tile, rubber, and granite. At the same time, I modified lighting angles. Because the textures and lights exist natively in the 3D world, the vision encoder is forced to learn true geometric edges and spatial depth across physically plausible surfaces.

After training across this material distribution for 14.7 million steps, the policy handles every table I throw at it. It succeeds on 99.9% of 1,022 evaluation episodes with randomized textures, including the woodgrain table that previously collapsed it to 5.7%. Woodgrain is one of the eight training textures though, so this is not a held-out test, and I have not tried textures outside the set.

Fig. 2 Table Texture Domain Randomization in Action

The hardened policy operating across shifting table materials. By training over a diverse set of 3D materials, the vision encoder learns to track the cube independently of background appearance.

Demonstrating visual invariance. As the table surface swaps dynamically across different materials, the hardened policy continues its reach, grasp, and placement sequence without hesitation.
Evaluation EnvironmentHome Success (1,022 to 1,024 episodes each)Approach / Lift / Place
Woodgrain Table (Nominal, in the training set)99.8%100% / 100% / 100%
Dark Table (Baseline Surface)100.0%100% / 100% / 100%
Random Multi-Texture (Active DR)99.9%100% / 100% / 100%

Realistic Movements and Some Final Reward Tuning

The policy is now more robust against environment changes, yee haw! However, it can still fail on physical hardware if it commands unrealistic motor movements so our work is not done yet.

In the simulation, the neural network can output drastic target changes from one step to the next. A poorly simulated motor will obediently attempt to snap to that position in a single millisecond tick. In the real world though, electric motors have rotor inertia, winding inductance, and signal latency. Demanding instantaneous torque steps causes current spikes, gearbox chatter, and controller instability. It could even break your robot, which would make all stakeholders very sad.

Modeling Motor Delays and Slew Limits

To prevent the network from relying on impossible acceleration, I added two dynamic filters. First, a discrete first-order low-pass filter for the setpoints commanded by the policy. This added roughly a 20 ms physical response lag at 50 Hz so the policy has to plan its trajectories knowing the arm takes time to catch up. Second, I added slew-rate limiting to the actuators. Now the gripper fingers move at a realistic physical velocity limit, taking roughly 0.42 seconds to complete a full stroke. Same concept applies for the joint actuators.

Reward Tweaking Continues Until Morale Improves

The Hovering Fix I Promised in Part I

Part I ended on a policy that learned to hover a few centimeters above the cube and pinch at air. The cause was the reward, not the network. The approach bonuses paid out at 15, 8, and 4 cm from the cube and nothing below that, but a real grasp needs the gripper within about 2 cm. The policy was closing its fingers at a median distance of 4.5 to 5 cm, right at the last payout, because getting any closer earned nothing.

The fix was a one-time +15 bonus for getting the gripper within 2 cm of the cube center, which is where the fingers actually straddle it. In a matched A/B from the same checkpoint and the same 200 updates, home-position success went from 0.699 to 0.889. Continuing for 600 more updates took it to 0.987, which is the 98.7% from Part I. At 200 updates the pinch statistics actually looked worse than the control, and it took the longer run for the median closing distance to drop from 5.3 cm to 1.7 cm, so I'd call the mechanism confirmed by the long run rather than by the A/B.

When I inspected eval footage of the policy after sorting the cube, I noticed a few strange behaviors. Two main issues you probably saw in the Part I videos were the little flick or toss the robot did to get a cube in the landing zone, and the second was that its wrist twisted at weird angles after completing the task.

This was caused by reward under-specification. The original retreat reward only measured the Cartesian distance from the gripper tip to the center of the workspace (a 10 cm sphere in space). The reward did not care how the arm's 7 joints were oriented, as long as the fingers ended up inside that sphere. Because twisting the wrist was slightly faster than retracting the whole arm, the optimizer happily exploited that shortcut.

To fix the odd wrist behavior, I replaced the simple Cartesian distance reward function I had with a combined Cartesian and joint-space potential one:

Combined retreat potential. The Cartesian term d_home pulls the end-effector home, while the joint-space term d_q penalizes deviation from the nominal ready posture q_ready across all 7 joints.

In this formulation, measures the physical distance from the gripper to the home position, while calculates the root-mean-square error between the robot's current 7 joint angles () and their canonical ready pose (). This creates a smooth gradient pulling the entire arm back into a clean, natural posture.

The earlier changes to the fingertip rubber modeling helped the toss problem along with a fix for the unconstrained release behavior. To teach it to set the cube down gently rather than flinging it, I added a precision placement bonus (awarding up to +8.0 scaled by how close the cube lands to dead center) and a gentle touchdown bonus (up to +6.0, only paid out if vertical touchdown speed is under 5 cm/s). Paired with penalties for high descent speed and an instant failure if the cube bounces or slides away after release, the policy learned that a deliberate, cushioned touchdown was vastly more profitable than a lucky throw.

All of these changes come together to make the overall robot behavior so much smoother and better! The final policy reaches 98.0% home success with all of them on, a bit below the 99.9% of the texture run. In fact, it looked so good that I had to do the robot justice and snazz the recordings of it up. I recorded the trajectories of the robot in MuJoCo and imported them into Blender. That is why these final videos look so much nicer. The movement of the robot is still 100% coming from the trained policy, I just made it more visually appealing to look at.

Fig. 3 Putting It All Together. Hardened Rollout and Tuned Rewards.

The policy operating under realistic actuator lag, lighting randomization, and joint-space posture shaping. Motions are smooth and continuous, and the arm settles into a clean ready configuration after placement.

A complete rollout under realistic motor dynamics, lighting randomization, and joint-space posture control. The arm moves smoothly, grips securely, and returns to a clean ready stance.

Where Things Stand

This project has been a painful reminder that reinforcement learning models are exceptional at finding shortcuts in simulation and reward policies. If you leave contacts frictionless, actuators instantaneous, backgrounds static, or rewards unconstrained, the policy will exploit those simplifications every time. It takes quite a bit of back and forth fiddling to get it right.

This was a massive step forward for this project. Now I can train on the cloud, I have policies that are becoming more robust, and I have a Blender pipeline setup to capture beautiful renders of the new and improved robot. So what is next? I could try to actually take it sim-to-real, but I don't have a spare $25,000-$40,000 USD lying around to buy this robot arm unfortunately. So, I will continue to make the things it does in the simulation more and more applicable to the real world! Next up is a harder task: picking items off a moving conveyor belt and stacking them in bins. It will take a while, so I'll post an update once I have findings. See you next time.

Have feedback or thoughts?

Found a typo, have a question on the derivations, or want to discuss this topic?

Send Feedback →
RoboticsReinforcement LearningMuJoCoCloud ComputingContact PhysicsDomain Randomization