Skip to content

Commit

Permalink
dora + model interchangeablility
Browse files Browse the repository at this point in the history
  • Loading branch information
nathanjzhao committed Aug 16, 2024
1 parent 9a6b6da commit 1d2a62c
Show file tree
Hide file tree
Showing 150 changed files with 2,346 additions and 100,158 deletions.
5 changes: 4 additions & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -77,4 +77,7 @@ jspm_packages/
#!node_modules/three/examples/jsm/utils/
#!node_modules/three/examples/jsm/utils/*
#!node_modules/three/examples/jsm/loaders/
#!node_modules/three/examples/jsm/loaders/*
#!node_modules/three/examples/jsm/loaders/*

scratch/
models/
73 changes: 2 additions & 71 deletions README.md
Original file line number Diff line number Diff line change
@@ -1,72 +1,3 @@
<p align="center">
<a href="https://zalo.github.io/mujoco_wasm/"><img src="./examples/MuJoCoWasmLogo.png" href></a>
</p>
<p align="left">
<a href="https://github.com/zalo/mujoco_wasm/deployments/activity_log?environment=github-pages">
<img src="https://img.shields.io/github/deployments/zalo/mujoco_wasm/github-pages?label=Github%20Pages%20Deployment" title="Github Pages Deployment"></a>
<!--<a href="https://github.com/zalo/mujoco_wasm/deployments/activity_log?environment=Production">
<img src="https://img.shields.io/github/deployments/zalo/mujoco_wasm/Production?label=Vercel%20Deployment" title="Vercel Deployment"></a> -->
<!--<a href="https://lgtm.com/projects/g/zalo/mujoco_wasm/context:javascript">
<img alt="Language grade: JavaScript" src="https://img.shields.io/lgtm/grade/javascript/g/zalo/mujoco_wasm.svg?logo=lgtm&logoWidth=18"/></a> -->
<a href="https://github.com/zalo/mujoco_wasm/commits/main">
<img src="https://img.shields.io/github/last-commit/zalo/mujoco_wasm" title="Last Commit Date"></a>
<a href="https://github.com/zalo/mujoco_wasm/blob/main/LICENSE">
<img src="https://img.shields.io/badge/license-MIT-brightgreen" title="License: MIT"></a>
</p>
# Interactive MUJOCO

## The Power of MuJoCo in your Browser.

Load and Run MuJoCo 2.3.1 Models using JavaScript and WebAssembly.

This repo is a fork of @stillonearth 's starter repository, adding tons of functionality and a comprehensive example scene.

### [See the Live Demo Here](https://zalo.github.io/mujoco_wasm/)

### [See a more Advanced Example Here](https://kzakka.com/robopianist/)

## Building

**1. Install emscripten**

**2. Build the mujoco_wasm Binary**

On Linux, use:
```bash
mkdir build
cd build
emcmake cmake ..
make
```

On Windows, run `build_windows.bat`.

*3. (Optional) Update MuJoCo libs*

Build MuJoCo libs with wasm target and place to lib. Currently v2.3.1 included.

## JavaScript API

```javascript
import load_mujoco from "./mujoco_wasm.js";

// Load the MuJoCo Module
const mujoco = await load_mujoco();

// Set up Emscripten's Virtual File System
mujoco.FS.mkdir('/working');
mujoco.FS.mount(mujoco.MEMFS, { root: '.' }, '/working');
mujoco.FS.writeFile("/working/humanoid.xml", await (await fetch("./examples/scenes/humanoid.xml")).text());

// Load in the state from XML
let model = new mujoco.Model("/working/humanoid.xml");
let state = new mujoco.State(model);
let simulation = new mujoco.Simulation(model, state);
```

Typescript definitions are available.

## Work In Progress Disclaimer

So far, most mjModel and mjData state variables and functions (that do not require custom structs) are exposed.

At some point, I'd like to de-opinionate the binding and make it match the original MuJoCo API better.
When building new models, make sure to run `generate_index.py`
95 changes: 87 additions & 8 deletions examples/main.js
Original file line number Diff line number Diff line change
Expand Up @@ -25,7 +25,7 @@ export class MuJoCoDemo {
this.simulation = new mujoco.Simulation(this.model, this.state);

// Define Random State Variables
this.params = { scene: initialScene, paused: false, help: false, ctrlnoiserate: 0.0, ctrlnoisestd: 0.0, keyframeNumber: 0 };
this.params = { scene: initialScene, paused: false, useModel: true, help: false, ctrlnoiserate: 0.0, ctrlnoisestd: 0.0, keyframeNumber: 0 };
this.mujoco_time = 0.0;
this.bodies = {}, this.lights = {};
this.tmpVec = new THREE.Vector3();
Expand Down Expand Up @@ -70,6 +70,7 @@ export class MuJoCoDemo {

this.actuatorNames = [];
this.actuatorRanges = [];
this.loadPPOModel();

window.addEventListener('resize', this.onWindowResize.bind(this));

Expand All @@ -89,10 +90,10 @@ export class MuJoCoDemo {
this.gui = new GUI();
setupGUI(this);

this.initializeActuators();
// this.initializeActuators();
}


// does this ordering align with the ordering of the model output?
initializeActuators() {
const textDecoder = new TextDecoder();
for (let i = 0; i < this.model.nu; i++) {
Expand All @@ -109,6 +110,53 @@ export class MuJoCoDemo {
}


// should re-get pawel-diff
async loadPPOModel() {
switch (this.params.scene) {
case 'humanoid.xml':
this.ppo_model = await tf.loadLayersModel('models/2_frame/model.json');
this.getObservation = () => this.getObservationSkeleton(2, 10, 6);
break;
case 'blank':
this.ppo_model = await tf.loadLayersModel('models/cvals+2_frames/model.json');
break;
case 'brax_humanoid.xml':
this.ppo_model = await tf.loadLayersModel('models/brax_humanoid_cvalless_just_stand/model.json');
this.getObservation = () => this.getObservationSkeleton(0, -1, -1);
break;
case 'brax_humanoidstandup.xml':
this.ppo_model = await tf.loadLayersModel('models/brax_humanoid_standup/model.json');
this.getObservation = () => this.getObservationSkeleton(0, 20, 12);
break;
default:
throw new Error(`Unknown model path: ${this.params.scene}`);
}
}

getObservationSkeleton(qpos_slice, cinert_slice, cvel_slice) {
const qpos = this.simulation.qpos.slice(qpos_slice);
const qvel = this.simulation.qvel;
const cinert = cinert_slice !== -1 ? this.simulation.cinert.slice(cinert_slice) : [];
const cvel = cvel_slice !== -1 ? this.simulation.cvel.slice(cvel_slice) : [];
const qfrc_actuator = this.simulation.qfrc_actuator;

// console.log('qpos length:', qpos.length);
// console.log('qvel length:', qvel.length);
// console.log('cinert length:', cinert.length);
// console.log('cvel length:', cvel.length);
// console.log('qfrc_actuator length:', qfrc_actuator.length);

const obsComponents = [
...qpos,
...qvel,
...cinert,
...cvel,
...qfrc_actuator
];

return obsComponents;
}

handleKeyPress(event) {
const key = event.key.toLowerCase();
const stepSize = 0.1;
Expand Down Expand Up @@ -184,7 +232,41 @@ export class MuJoCoDemo {
this.controls.update();

if (!this.params["paused"]) {
// this.simulation.xfrc_applied.fill(3.0)

if (this.ppo_model && this.params["useModel"]) {
const observationArray = this.getObservation();
const inputTensor = tf.tensor2d([observationArray]);
const resultTensor = this.ppo_model.predict(inputTensor);

resultTensor.data().then(data => {
// console.log('Model output:', data);

// Assuming the model output corresponds to actuator values
for (let i = 0; i < data.length; i++) {
// Ensure the actuator index is within bounds
if (i < this.simulation.ctrl.length) {

let clippedValue = Math.max(-1, Math.min(1, data[i]));

let [min, max] = this.actuatorRanges[i];

// Scale to fit between min and max
let newValue = min + (clippedValue + 1) * (max - min) / 2;

// Update the actuator value
this.simulation.ctrl[i] = newValue;

// Optionally, update the corresponding parameter
this.params[this.actuatorNames[i]] = newValue;
} else {
console.error('Model output index out of bounds:', i);
}
}
});
}

// console.log(this.model)
// console.log(this.model.getOptions().timestep);

let timestep = this.model.getOptions().timestep;
if (timeMS - this.mujoco_time > 35.0) { this.mujoco_time = timeMS; }
Expand All @@ -199,12 +281,9 @@ export class MuJoCoDemo {
let currentCtrl = this.simulation.ctrl;
for (let i = 0; i < currentCtrl.length; i++) {
currentCtrl[i] = rate * currentCtrl[i] + scale * standardNormal();
this.params["Actuator " + i] = currentCtrl[i];
this.params[this.actuatorNames[i]] = currentCtrl[i];
}
}

// actions = this.simulation.ctrl
console.log("c", this.simulation.ctrl)

// Clear old perturbations, apply new ones.
for (let i = 0; i < this.simulation.qfrc_applied.length; i++) { this.simulation.qfrc_applied[i] = 0.0; }
Expand Down
Loading

0 comments on commit 1d2a62c

Please sign in to comment.