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

rollout use size_t for array size check and pointer arithmetic

Browse Source
This commit is contained in:
Levi Burner
2025-02-03 22:22:37 -05:00
parent 19624aef9d
commit df71467eed
1 changed files with 11 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/rollout.cc
+11 -9
View File
@@ -75,10 +75,11 @@ void _unsafe_rollout(std::vector<const mjModel*>& m, mjData* d, int start_roll,
const mjtNum* state0, const mjtNum* warmstart0, const mjtNum* state0, const mjtNum* warmstart0,
const mjtNum* control, mjtNum* state, mjtNum* sensordata) { const mjtNum* control, mjtNum* state, mjtNum* sensordata) {
// sizes // sizes
int nstate = mj_stateSize(m[0], mjSTATE_FULLPHYSICS); size_t nstate = static_cast<size_t>(mj_stateSize(m[0], mjSTATE_FULLPHYSICS));
int ncontrol = mj_stateSize(m[0], control_spec); size_t ncontrol = static_cast<size_t>(mj_stateSize(m[0], control_spec));
int nv = m[0]->nv, nbody = m[0]->nbody, neq = m[0]->neq; size_t nv = static_cast<size_t>(m[0]->nv);
int nsensordata = m[0]->nsensordata; int nbody = m[0]->nbody, neq = m[0]->neq;
size_t nsensordata = static_cast<size_t>(m[0]->nsensordata);
// clear user inputs if unspecified // clear user inputs if unspecified
if (!(control_spec & mjSTATE_CTRL)) { if (!(control_spec & mjSTATE_CTRL)) {
@@ -92,7 +93,7 @@ void _unsafe_rollout(std::vector<const mjModel*>& m, mjData* d, int start_roll,
} }
// loop over rollouts // loop over rollouts
for (int r = start_roll; r < end_roll; r++) { for (size_t r = start_roll; r < end_roll; r++) {
// clear user inputs if unspecified // clear user inputs if unspecified
if (!(control_spec & mjSTATE_MOCAP_POS)) { if (!(control_spec & mjSTATE_MOCAP_POS)) {
for (int i = 0; i < nbody; i++) { for (int i = 0; i < nbody; i++) {
@@ -128,7 +129,7 @@ void _unsafe_rollout(std::vector<const mjModel*>& m, mjData* d, int start_roll,
} }
// roll out trajectory // roll out trajectory
for (int t = 0; t < nstep; t++) { for (size_t t = 0; t < nstep; t++) {
// check for warnings // check for warnings
bool nwarning = false; bool nwarning = false;
for (int i = 0; i < mjNWARNING; i++) { for (int i = 0; i < mjNWARNING; i++) {
@@ -141,7 +142,7 @@ void _unsafe_rollout(std::vector<const mjModel*>& m, mjData* d, int start_roll,
// if any warnings, fill remaining outputs with current outputs, break // if any warnings, fill remaining outputs with current outputs, break
if (nwarning) { if (nwarning) {
for (; t < nstep; t++) { for (; t < nstep; t++) {
int step = r*nstep + t; size_t step = r*static_cast<size_t>(nstep) + t;
if (state) { if (state) {
mj_getState(m[r], d, state + step*nstate, mjSTATE_FULLPHYSICS); mj_getState(m[r], d, state + step*nstate, mjSTATE_FULLPHYSICS);
} }
@@ -152,7 +153,7 @@ void _unsafe_rollout(std::vector<const mjModel*>& m, mjData* d, int start_roll,
break; break;
} }
int step = r*nstep + t; size_t step = r*static_cast<size_t>(nstep) + t;
// controls // controls
if (control) { if (control) {
@@ -226,7 +227,8 @@ mjtNum* get_array_ptr(std::optional<const py::array_t<mjtNum>> arg,
py::buffer_info info = arg->request(); py::buffer_info info = arg->request();
// check size // check size
int expected_size = nbatch * nstep * dim; size_t expected_size =
static_cast<size_t>(nbatch) * static_cast<size_t>(nstep) * static_cast<size_t>(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;