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

clean up rollout.cc

Browse Source 
PiperOrigin-RevId: 599849967
Change-Id: I899b28bf78bf0414b07314c24cfe6e24441eba30
This commit is contained in:
Yuval Tassa
2024-01-19 09:02:27 -08:00
committed by Copybara-Service
parent e0864ab7f2
commit a23a368778
1 changed files with 15 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/rollout.cc
+15 -17
View File
@@ -12,15 +12,14 @@
// See the License for the specific language governing permissions and // See the License for the specific language governing permissions and
// limitations under the License. // limitations under the License.
#include <array>
#include <cstdio>
#include <iostream> #include <iostream>
#include <optional> #include <optional>
#include <sstream> #include <sstream>
#include <string>
#include "functions.h" #include <mujoco/mujoco.h>
#include "errors.h"
#include "raw.h" #include "raw.h"
#include "structs.h"
#include <pybind11/buffer_info.h> #include <pybind11/buffer_info.h>
#include <pybind11/numpy.h> #include <pybind11/numpy.h>
#include <pybind11/pybind11.h> #include <pybind11/pybind11.h>
@@ -74,7 +73,6 @@ void _unsafe_rollout(const mjModel* m, mjData* d, int nstate, int nstep,
// loop over initial states // loop over initial states
for (int s=0; s < nstate; s++) { for (int s=0; s < nstate; s++) {
// set initial state // set initial state
if (state0) { if (state0) {
mju_copy(d->qpos, state0 + s*nqva, nq); mju_copy(d->qpos, state0 + s*nqva, nq);
@@ -108,9 +106,9 @@ void _unsafe_rollout(const mjModel* m, mjData* d, int nstate, int nstep,
mju_zero(d->xfrc_applied, 6*nbody); mju_zero(d->xfrc_applied, 6*nbody);
} }
if (!mocap) { if (!mocap) {
for (int j=0; j<nbody; j++) { for (int j=0; j < nbody; j++) {
int id = m->body_mocapid[j]; int id = m->body_mocapid[j];
if (id>=0) { if (id >= 0) {
mju_copy3(d->mocap_pos+3*id, m->body_pos+3*j); mju_copy3(d->mocap_pos+3*id, m->body_pos+3*j);
mju_copy4(d->mocap_quat+4*id, m->body_quat+4*j); mju_copy4(d->mocap_quat+4*id, m->body_quat+4*j);
} }
@@ -174,7 +172,8 @@ mjtNum* get_array_ptr(std::optional<const py::array_t<mjtNum>> arg,
int expected_size = nstate * nstep * dim; int expected_size = nstate * nstep * dim;
if (info.size != expected_size) { if (info.size != expected_size) {
std::ostringstream msg; std::ostringstream msg;
msg << name << ".size should be " << expected_size << ", got " << info.size; msg << name << ".size should be " << expected_size <<
", got " << info.size;
throw py::value_error(msg.str()); throw py::value_error(msg.str());
} }
return static_cast<mjtNum*>(info.ptr); return static_cast<mjtNum*>(info.ptr);
@@ -200,7 +199,6 @@ PYBIND11_MODULE(_rollout, pymodule) {
std::optional<const PyCArray> state, std::optional<const PyCArray> state,
std::optional<const PyCArray> sensordata std::optional<const PyCArray> sensordata
) { ) {
const raw::MjModel* model = m.get(); const raw::MjModel* model = m.get();
raw::MjData* data = d.get(); raw::MjData* data = d.get();
@@ -213,7 +211,8 @@ PYBIND11_MODULE(_rollout, pymodule) {
int nqva = model->nq + model->nv + model->na; int nqva = model->nq + model->nv + model->na;
mjtNum* init_state_ptr = mjtNum* init_state_ptr =
get_array_ptr(init_state, "initial_state", nstate, 1, nqva); get_array_ptr(init_state, "initial_state", nstate, 1, nqva);
mjtNum* ctrl_ptr = get_array_ptr(ctrl, "ctrl", nstate, nstep, model->nu); mjtNum* ctrl_ptr =
get_array_ptr(ctrl, "ctrl", nstate, nstep, model->nu);
mjtNum* qfrc_ptr = mjtNum* qfrc_ptr =
get_array_ptr(qfrc, "qfrc_applied", nstate, nstep, model->nv); get_array_ptr(qfrc, "qfrc_applied", nstate, nstep, model->nv);
mjtNum* xfrc_ptr = mjtNum* xfrc_ptr =
@@ -222,11 +221,11 @@ PYBIND11_MODULE(_rollout, pymodule) {
get_array_ptr(mocap, "mocap", nstate, nstep, 7*model->nmocap); get_array_ptr(mocap, "mocap", nstate, nstep, 7*model->nmocap);
mjtNum* init_time_ptr = mjtNum* init_time_ptr =
get_array_ptr(init_time, "init_time", nstate, 1, 1); get_array_ptr(init_time, "init_time", nstate, 1, 1);
mjtNum* init_warmstart_ptr = mjtNum* init_warmstart_ptr = get_array_ptr(
get_array_ptr(init_warmstart, "init_warmstart", nstate, 1, model->nv); init_warmstart, "init_warmstart", nstate, 1, model->nv);
mjtNum* state_ptr = get_array_ptr(state, "state", nstate, nstep, nqva); mjtNum* state_ptr = get_array_ptr(state, "state", nstate, nstep, nqva);
mjtNum* sensordata_ptr = mjtNum* sensordata_ptr = get_array_ptr(sensordata, "sensordata", nstate,
get_array_ptr(sensordata, "sensordata", nstate, nstep, model->nsensordata); nstep, model->nsensordata);
// perform rollouts // perform rollouts
{ {
@@ -255,9 +254,8 @@ PYBIND11_MODULE(_rollout, pymodule) {
py::arg("sensordata") = py::none(), py::arg("sensordata") = py::none(),
py::doc(rollout_doc) py::doc(rollout_doc)
); );
}
} // namespace } // namespace
} } // namespace mujoco::python
}