Add Differentiable Physics: Mass-Spring System example #1332

AbhiLegend · 2025-04-29T13:04:12Z

Summary

This PR adds a simple Differentiable Physics example implementing a Mass-Spring System using PyTorch.

Particles are connected by springs and evolve under the influence of spring forces and gravity.
The system is fully differentiable, allowing optimization of particle positions via gradient descent.
Semi-implicit Euler integration is used for better numerical stability.
The example does not require any external dependencies beyond PyTorch.

Files Added

examples/differentiable_physics/mass_spring.py
examples/differentiable_physics/README.md
examples/differentiable_physics/requirements.txt

Features

Training mode: Optimize particle positions to match a target configuration.
Evaluation mode: Forward simulate without optimization.
Particle 0 is pinned at (0,0) to prevent system drift.
Gravity strength is customizable.
Fully differentiable simulation with simple, clean code structure.

Testing

Verified mass_spring.py runs successfully in both train and eval modes.
Added a bash function to run_python_examples.sh for CI testing (5 epochs, 3 steps for quick test).
Local testing was performed with Git Bash on Windows (python -m uv is used for compatibility).
Confirmed loss convergence and expected particle behavior.

Notes

Minor residual errors are expected due to the elasticity of springs and gravitational force.
Example is designed to prioritize educational clarity and maintain a low dependency footprint.

Thank you for reviewing!

netlify · 2025-04-29T13:04:36Z

✅ Deploy Preview for pytorch-examples-preview canceled.

Name	Link
🔨 Latest commit	`f1a806e`
🔍 Latest deploy log	https://app.netlify.com/projects/pytorch-examples-preview/deploys/682ce2476390c90008af215c

AbhiLegend · 2025-04-30T20:02:13Z

@msaroufim I have just updated the example with PR, request you to check if it fits in the "differentiable physics" section

msaroufim · 2025-05-13T21:49:38Z

I don't see any testing changes that were made, you probably want to update this file https://github.com/pytorch/examples/blob/main/run_python_examples.sh

I do also think it's worth making some visualization for the work and putting it in the main README

AbhiLegend · 2025-05-13T22:02:50Z

I don't see any testing changes that were made, you probably want to update this file https://github.com/pytorch/examples/blob/main/run_python_examples.sh

I do also think it's worth making some visualization for the work and putting it in the main README

I have updated the PR.Thanks

AbhiLegend · 2025-05-13T22:33:25Z

I don't see any testing changes that were made, you probably want to update this file https://github.com/pytorch/examples/blob/main/run_python_examples.sh

I do also think it's worth making some visualization for the work and putting it in the main README

Updated all the changes.

AbhiLegend · 2025-05-16T12:26:03Z

@msaroufim , can you please review the PR? Please lemme know if something else is required for this example. Thanks :)

jafraustro · 2025-05-18T04:27:43Z

differentiable_physics/requirements.txt

@@ -0,0 +1 @@
+torch


Hi @AbhiLegend,

matplotlib dependency is missing.

jafraustro · 2025-05-18T04:34:13Z

differentiable_physics/mass_spring.py

+
+
+def train(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


It would be better to use accelerator API here

Suggested change

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

device = torch.accelerator.current_accelerator() if torch.accelerator.is_available() else torch.device("cpu")

print(f"Using device: {device}")

jafraustro · 2025-05-18T04:35:01Z

differentiable_physics/readme.md

+
+- Python 3.8+
+- PyTorch
+


matplotlib is missing

Add

Suggested change

pip install -r requirements.txt

to be consistent with other examples

jafraustro · 2025-05-18T04:37:41Z

differentiable_physics/mass_spring.py

+
+
+def evaluate(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


Suggested change

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

device = torch.accelerator.current_accelerator() if torch.accelerator.is_available() else torch.device("cpu")

print(f"Using device: {device}")

@jafraustro Final update:

Added accelerator code to mass_spring.py

requirements.txt updated with matplotlib dependency
Ready for review

Hi, I do not see matplotlib in the requirements.txt file

https://github.com/AbhiLegend/examples/blob/8a7cf5e4fcc06582a968b7ba8348e99c9faa5b27/differentiable_physics/requirements.txt

AbhiLegend · 2025-05-18T21:21:22Z

@jafraustro I made the changes.

jafraustro

LGTM

dvrogozh · 2025-05-19T18:33:02Z

run_python_examples.sh

 #
-# To test examples on CUDA accelerator, run as:
-#   USE_CUDA=True ./run_python_examples.sh
+# To run specific examples:


I think these changes be better moved to separate PR. These are unrelated to the example being added and some are arguable.

Made it to as it was.

dvrogozh · 2025-05-19T18:34:39Z

run_python_examples.sh

 #
-# To test examples on hardware accelerator (CUDA, MPS, XPU, etc.), run as:
-#   USE_ACCEL=True ./run_python_examples.sh
-# NOTE: USE_ACCEL relies on torch.accelerator API and not all examples are converted


Why a note was dropped? it highlights the planned work which is not yet done and I believe is helpful for maintenance.

dvrogozh · 2025-05-19T18:35:27Z

run_python_examples.sh

 #
-# Script requires uv to be installed. When executed, script will install prerequisites from


This highlights the convention that each example is expected to have requirements.txt and discusses the behavior. Can this, please, be restored?

restored to original

dvrogozh · 2025-05-19T18:37:55Z

run_python_examples.sh

 #
-# To force script to create virtual environment for each example, run as:
+# To force venv per example:


venv might be confused with the type of virtual environment. I would prefer to use full naming "virtual environment per example".

dvrogozh · 2025-05-19T18:38:31Z

run_python_examples.sh

 #   VIRTUAL_ENV=".venv" ./run_python_examples.sh
-# Script will remove environments it creates in a teardown step after execution of each example.


Again, I don't see why this needs to be dropped. That's explicit clarification of the behavior. It's helpful.

dvrogozh · 2025-05-19T18:38:56Z

run_python_examples.sh


 BASE_DIR="$(pwd)/$(dirname $0)"
 source $BASE_DIR/utils.sh

-# TODO: Leave only USE_ACCEL and drop USE_CUDA once all examples will be converted


This work is not completed. Why drop the TODO?

Made the changes.

dvrogozh · 2025-05-19T18:41:08Z

run_python_examples.sh

  echo "training..."
  uv run train.py --epochs 1 --dev_every 1 --no-bidirectional --dry-run || error "couldn't train snli"
 }

 function fx() {
-  # uv run custom_tracer.py || error "fx custom tracer has failed" UnboundLocalError: local variable 'tabulate' referenced before assignment


cleanup, here and in few places below. But someone needs to clarify whether this can be dropped or not. In any case - better not to mix up such changes with the new example.

dvrogozh · 2025-05-19T18:45:18Z

run_python_examples.sh

+  pushd differentiable_physics
+  python -m uv run mass_spring.py --mode train --epochs 5 --steps 3 || error "differentiable_physics example failed"
+  popd


What's the reason to use different pattern for tests just for this example? One potential problem - python might not actually exist on the system as typically you get python3 and need to create a symbolic link to it in one way or another.

Suggested change

pushd differentiable_physics

python -m uv run mass_spring.py --mode train --epochs 5 --steps 3 || error "differentiable_physics example failed"

popd

uv run mass_spring.py --mode train --epochs 5 --steps 3 || error "differentiable_physics example failed"

Rectified it.

AbhiLegend · 2025-05-19T22:10:43Z

@dvrogozh made the changes as perrequest.

jafraustro · 2025-05-19T23:37:58Z

run_python_examples.sh

  exit 1
-fi
+
+fi


add the empty line at the end

dvrogozh · 2025-05-20T17:57:01Z

run_python_examples.sh

@@ -154,6 +154,9 @@ function vision_transformer() {

 function word_language_model() {
  uv run main.py --epochs 1 --dry-run $CUDA_FLAG --mps || error "word_language_model failed"
+  for model in "RNN_TANH" "RNN_RELU" "LSTM" "GRU" "Transformer"; do


This change is unrelated with the PR. Remove. Or just rebase on top of latest main.

dvrogozh · 2025-05-20T17:57:52Z

differentiable_physics/requirements.txt

@@ -0,0 +1,3 @@
+torch


Due to usage of torch.accelerator:

Suggested change

torch

torch>=2.6

AbhiLegend · 2025-05-20T20:17:05Z

@dvrogozh @jafraustro I have made the changes.

Add Differentiable Physics: Mass-Spring System example

7d3d2a1

facebook-github-bot added the cla signed label Apr 29, 2025

AbhiLegend mentioned this pull request May 1, 2025

New examples requested #1131

Open

12 tasks

Add differentiable_physics to run_all() in test script

39a0c8e

AbhiLegend requested a review from msaroufim as a code owner May 13, 2025 22:00

AbhiLegend added 2 commits May 14, 2025 03:40

Resolved merge conflict and integrated differentiable_physics test

3d01b48

Add visualization and update training code in mass_spring.py

77c8abc

jafraustro reviewed May 18, 2025

View reviewed changes

AbhiLegend added 4 commits May 18, 2025 19:01

Finalize differentiable_physics with visualization and CI integration

84072d5

Finalize differentiable_physics with visualization and CI integration

8a7cf5e

Finalize differentiable_physics with the updates

f20fa20

Update requirements.txt for differentiable_physics

e14251f

jafraustro approved these changes May 19, 2025

View reviewed changes

dvrogozh suggested changes May 19, 2025

View reviewed changes

Update run_python_examples.sh to test differentiable_physics in CI

96e04f1

jafraustro reviewed May 19, 2025

View reviewed changes

run_python_examples.sh Outdated

exit 1

fi

fi

Copy link

Contributor

jafraustro May 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

add the empty line at the end

dvrogozh reviewed May 20, 2025

View reviewed changes

AbhiLegend added 2 commits May 21, 2025 01:35

Add mass spring example and update requirements

35b0afa

Add mass spring example and update requirements

f1a806e



		def train(args):
		device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	device = torch.accelerator.current_accelerator() if torch.accelerator.is_available() else torch.device("cpu")
	print(f"Using device: {device}")



		def evaluate(args):
		device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

		#
		# Script requires uv to be installed. When executed, script will install prerequisites from

		# VIRTUAL_ENV=".venv" ./run_python_examples.sh
		# Script will remove environments it creates in a teardown step after execution of each example.

                 exit 1
-              fi
+              fi

Add Differentiable Physics: Mass-Spring System example #1332

Are you sure you want to change the base?

Add Differentiable Physics: Mass-Spring System example #1332

Uh oh!

Conversation

AbhiLegend commented Apr 29, 2025

Summary

Files Added

Features

Testing

Notes

Uh oh!

netlify bot commented Apr 29, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅ Deploy Preview for pytorch-examples-preview canceled.

Uh oh!

AbhiLegend commented Apr 30, 2025

Uh oh!

msaroufim commented May 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

AbhiLegend commented May 13, 2025

Uh oh!

AbhiLegend commented May 13, 2025

Uh oh!

AbhiLegend commented May 16, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AbhiLegend commented May 18, 2025

Uh oh!

jafraustro left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AbhiLegend commented May 19, 2025

Uh oh!

netlify bot commented Apr 29, 2025 •

edited

Loading

msaroufim commented May 13, 2025 •

edited

Loading