chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add sensors for potential and kinetic energy.

Browse Source 
PiperOrigin-RevId: 716775375
Change-Id: Ic8ab7f1a51df970ab551dabe7c797cde13cd97fd
This commit is contained in:
Yuval Tassa
2025-01-17 13:09:29 -08:00
committed by Copybara-Service
parent 688d367733
commit 240a7afdee
15 changed files with 264 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/engine/engine_forward.c
+57 -5
View File
@@ -1013,6 +1013,38 @@ void mj_implicit(const mjModel* m, mjData* d) {
// return 1 if potential energy was computed by sensor, 0 otherwise
static int energyPosSensor(const mjModel* m) {
if (mjDISABLED(mjDSBL_SENSOR)) {
return 0;
}
for (int i=0; i < m->nsensor; i++) {
if (m->sensor_type[i] == mjSENS_E_POTENTIAL) {
return 1;
}
}
return 0;
}
// return 1 if kinetic energy was computed by sensor, 0 otherwise
static int energyVelSensor(const mjModel* m) {
if (mjDISABLED(mjDSBL_SENSOR)) {
return 0;
}
for (int i=0; i < m->nsensor; i++) {
if (m->sensor_type[i] == mjSENS_E_KINETIC) {
return 1;
}
}
return 0;
}
//-------------------------- top-level API ---------------------------------------------------------
// forward dynamics with skip; skipstage is mjtStage
@@ -1022,21 +1054,33 @@ void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor)
// position-dependent
if (skipstage < mjSTAGE_POS) {
mj_fwdPosition(m, d);
int energyPos = 0;
if (!skipsensor) {
mj_sensorPos(m, d);
energyPos = energyPosSensor(m);
}
if (mjENABLED(mjENBL_ENERGY)) {
mj_energyPos(m, d);
if (!energyPos) {
if (mjENABLED(mjENBL_ENERGY)) {
mj_energyPos(m, d);
} else {
d->energy[0] = d->energy[1] = 0;
}
}
}
// velocity-dependent
if (skipstage < mjSTAGE_VEL) {
mj_fwdVelocity(m, d);
int energyVel = 0;
if (!skipsensor) {
mj_sensorVel(m, d);
energyVel = energyVelSensor(m);
}
if (mjENABLED(mjENBL_ENERGY)) {
if (mjENABLED(mjENBL_ENERGY) && !energyVel) {
mj_energyVel(m, d);
}
}
@@ -1111,10 +1155,18 @@ void mj_step1(const mjModel* m, mjData* d) {
mj_checkVel(m, d);
mj_fwdPosition(m, d);
mj_sensorPos(m, d);
mj_energyPos(m, d);
if (!energyPosSensor(m)) {
if (mjENABLED(mjENBL_ENERGY)) {
mj_energyPos(m, d);
} else {
d->energy[0] = d->energy[1] = 0;
}
}
mj_fwdVelocity(m, d);
mj_sensorVel(m, d);
mj_energyVel(m, d);
if (mjENABLED(mjENBL_ENERGY) && !energyVelSensor(m)) {
mj_energyVel(m, d);
}
if (mjcb_control) {
mjcb_control(m, d);
}
src/engine/engine_io.c
+2
View File
@@ -2056,6 +2056,8 @@ static int sensorSize(mjtSensor sensor_type, int sensor_dim) {
case mjSENS_TENDONLIMITVEL:
case mjSENS_TENDONLIMITFRC:
case mjSENS_GEOMDIST:
case mjSENS_E_POTENTIAL:
case mjSENS_E_KINETIC:
case mjSENS_CLOCK:
return 1;
src/engine/engine_sensor.c
+10 -11
View File
@@ -458,6 +458,16 @@ void mj_sensorPos(const mjModel* m, mjData* d) {
}
break;
case mjSENS_E_POTENTIAL: // potential energy
mj_energyPos(m, d);
d->sensordata[adr] = d->energy[0];
break;
case mjSENS_E_KINETIC: // kinetic energy
mj_energyVel(m, d);
d->sensordata[adr] = d->energy[1];
break;
case mjSENS_CLOCK: // clock
d->sensordata[adr] = d->time;
break;
@@ -902,12 +912,6 @@ void mj_energyPos(const mjModel* m, mjData* d) {
int padr;
mjtNum dif[3], quat[4], stiffness;
// disabled: clear and return
if (!mjENABLED(mjENBL_ENERGY)) {
d->energy[0] = d->energy[1] = 0;
return;
}
// init potential energy: -sum_i body(i).mass * mju_dot(body(i).pos, gravity)
d->energy[0] = 0;
if (!mjDISABLED(mjDSBL_GRAVITY)) {
@@ -996,11 +1000,6 @@ void mj_energyPos(const mjModel* m, mjData* d) {
// velocity-dependent energy (kinetic)
void mj_energyVel(const mjModel* m, mjData* d) {
// return if disabled (already cleared in potential)
if (!mjENABLED(mjENBL_ENERGY)) {
return;
}
mj_markStack(d);
mjtNum *vec = mjSTACKALLOC(d, m->nv, mjtNum);
src/user/user_objects.cc
+7 -1
View File
@@ -6206,7 +6206,11 @@ void mjCSensor::ResolveReferences(const mjCModel* m) {
((mjCGeom*)obj)->SetNotVisual();
}
} else if (type != mjSENS_CLOCK && type != mjSENS_PLUGIN && type != mjSENS_USER) {
} else if (type != mjSENS_E_POTENTIAL &&
type != mjSENS_E_KINETIC &&
type != mjSENS_CLOCK &&
type != mjSENS_PLUGIN &&
type != mjSENS_USER) {
throw mjCError(this, "invalid type in sensor");
}
@@ -6539,6 +6543,8 @@ void mjCSensor::Compile(void) {
}
break;
case mjSENS_E_POTENTIAL:
case mjSENS_E_KINETIC:
case mjSENS_CLOCK:
dim = 1;
needstage = mjSTAGE_POS;
src/xml/xml_native_reader.cc
+9 -1
View File
@@ -485,6 +485,8 @@ const char* MJCF[nMJCF][mjXATTRNUM] = {
{"distance", "*", "8", "name", "geom1", "geom2", "body1", "body2", "cutoff", "noise", "user"},
{"normal", "*", "8", "name", "geom1", "geom2", "body1", "body2", "cutoff", "noise", "user"},
{"fromto", "*", "8", "name", "geom1", "geom2", "body1", "body2", "cutoff", "noise", "user"},
{"e_potential", "*", "4", "name", "cutoff", "noise", "user"},
{"e_kinetic", "*", "4", "name", "cutoff", "noise", "user"},
{"clock", "*", "4", "name", "cutoff", "noise", "user"},
{"user", "*", "9", "name", "objtype", "objname", "datatype", "needstage",
"dim", "cutoff", "noise", "user"},
@@ -4172,7 +4174,13 @@ void mjXReader::Sensor(XMLElement* section) {
}
// global sensors
else if (type=="clock") {
else if (type=="e_potential") {
sensor->type = mjSENS_E_POTENTIAL;
sensor->objtype = mjOBJ_UNKNOWN;
} else if (type=="e_kinetic") {
sensor->type = mjSENS_E_KINETIC;
sensor->objtype = mjOBJ_UNKNOWN;
} else if (type=="clock") {
sensor->type = mjSENS_CLOCK;
sensor->objtype = mjOBJ_UNKNOWN;
}
src/xml/xml_native_reader.h
+1 -1
View File
@@ -101,7 +101,7 @@ class mjXReader : public mjXBase {
};
// MJCF schema
#define nMJCF 237
#define nMJCF 239
extern const char* MJCF[nMJCF][mjXATTRNUM];
#endif // MUJOCO_SRC_XML_XML_NATIVE_READER_H_
src/xml/xml_native_writer.cc
+6
View File
@@ -2156,6 +2156,12 @@ void mjXWriter::Sensor(XMLElement* root) {
break;
// global sensors
case mjSENS_E_POTENTIAL:
elem = InsertEnd(section, "potential");
break;
case mjSENS_E_KINETIC:
elem = InsertEnd(section, "kinetic");
break;
case mjSENS_CLOCK:
elem = InsertEnd(section, "clock");
break;