Predictive control on small scale robot (fruitfly)

Hello, I am trying to get [Fruitfly](https://github.com/google-deepmind/mujoco_menagerie/tree/main/flybody) model working with model predictive control but I am constantly getting `Rollout divergence at step` error. I am using the following set of model patch as well as task description.

fruitfly.xml.patch
```
diff --git a/flybody/fruitfly.xml b/flybody/fruitfly.xml
index e92beb0..393b2f9 100644
--- a/flybody/fruitfly.xml
+++ b/flybody/fruitfly.xml
@@ -1,7 +1,7 @@



-


@@ -882,7 +882,7 @@



-
+


-
+






-
+


-
+

```

task.xml
```

```

flybody.cc
```
#include "mjpc/tasks/flybody/flybody.h"

#include

#include
#include "mjpc/task.h"
#include "mjpc/utilities.h"

namespace mjpc {
std::string Flybody::XmlPath() const {
return GetModelPath("flybody/task.xml");
}
std::string Flybody::Name() const { return "Flybody"; }

// --------------- Residuals for quadrotor task ---------------
// Number of residuals: 4
// Residual (0): position - goal position
// Residual (1): linear velocity - goal linear velocity
// Residual (2): angular velocity - goal angular velocity
// Residual (3): control - hover control
// Number of parameters: 6
// Residuals are numbers that the MPC taretes
// ------------------------------------------------------------
void Flybody::ResidualFn::Residual(const mjModel* model, const mjData* data,
double* residuals) const {

// ---------- Residual (0) ----------
double* position = SensorByName(model, data, "position");
mju_sub(residuals, position, data->mocap_pos, 3); //distance from goal pos

// ---------- Residual (1) ----------
double* linear_velocity = SensorByName(model, data, "linear_velocity");
mju_copy(residuals + 3, linear_velocity, 3);

// ---------- Residual (2) ----------
double* angular_velocity = SensorByName(model, data, "angular_velocity");
mju_copy(residuals + 6, angular_velocity, 3);

// ---------- Residual (3) ----------
double thrust = (model->body_mass[0] + model->body_mass[1]) *
mju_norm3(model->opt.gravity) / model->nu;
for (int i = 0; i < model->nu; i++) {
residuals[9 + i] = data->ctrl[i] - thrust;
}
}

// ----- Transition for quadrotor task -----
void Flybody::TransitionLocked(mjModel* model, mjData* data) {
// set mode to GUI selection
if (mode > 0) {
current_mode_ = mode - 1;
} else {
// goal position
const double* goal_position = data->mocap_pos;

// system's position
double* position = SensorByName(model, data, "position");

// position error
double position_error[3];
mju_sub3(position_error, position, goal_position);
double position_error_norm = mju_norm3(position_error);

if (position_error_norm <= 5.0e-1) {
// update task state
current_mode_ += 1;
if (current_mode_ == model->nkey) {
current_mode_ = 0;
}
}
}

//std::printf("mocap_pos : %.2f", *(data->mocap_pos));
// set goal
//mju_copy3(data->mocap_pos, model->key_mpos + 3 * current_mode_);
//std::printf("mocap_pos : %.2f", *(data->mocap_pos));
//mju_copy4(data->mocap_quat, model->key_mquat + 4 * current_mode_);
}

} // namespace mjpc

```

fruitfly.h
```
#ifndef MJPC_TASKS_FLYBODY_FLYBODY_H_
#define MJPC_TASKS_FLYBODY_FLYBODY_H_

#include
#include
#include
#include "mjpc/task.h"

namespace mjpc {
class Flybody : public Task {
public:
std::string Name() const override;
std::string XmlPath() const override;
class ResidualFn : public mjpc::BaseResidualFn {
public:
ResidualFn(const ResidualFn& residual) = default;
explicit ResidualFn(const Flybody* task) : mjpc::BaseResidualFn(task) {}
// --------------- Residuals for quadrotor task ---------------
// Number of residuals: 5
// Residual (0): position - goal position
// Residual (1): orientation - goal orientation
// Residual (2): linear velocity - goal linear velocity
// Residual (3): angular velocity - goal angular velocity
// Residual (4): control
// Number of parameters: 6
// ------------------------------------------------------------
void Residual(const mjModel* model, const mjData* data,
double* residual) const override;
};

Flybody() : residual_(this) {}
void TransitionLocked(mjModel* model, mjData* data) override;

protected:
std::unique_ptr ResidualLocked() const override {
return std::make_unique(this);
}
ResidualFn* InternalResidual() override { return &residual_; }

private:
int current_mode_ = 0;
ResidualFn residual_;
};
} // namespace mjpc

#endif // MJPC_TASKS_FLYBODY_FLYBODY_H_
```

Essentially I am removing all actuators that are not wing and using the same task class as the quadrotor. As for current setup, without setting the horizon to below 0.02 and setting time to 0.1, the rollout fails completely and the agent does not move at all. With small horizon, the fly moves but does not fly or make reasonable consistent progress towards the goal.

My question is, is there a parameter in the mjpc system that is used specifically for small robots with actuator that produces large force in proportion to its mass?