Refactor the projectile launcher into a plugin.
Add a simple GUI for it that allows users to enable/disable the key binding as well as control some parameters. PiperOrigin-RevId: 872816645 Change-Id: I8abb2803e17b4f6195d447220c436d4701b24db8
This commit is contained in:
Haroon Qureshi
committed by
Copybara-Service
Copybara-Service
parent
28ad603f6b
commit
0083cf7fa8
@@ -0,0 +1,214 @@
|
||||
// Copyright 2026 DeepMind Technologies Limited
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include <random>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include <imgui.h>
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "experimental/platform/imgui_widgets.h"
|
||||
#include "experimental/platform/plugin.h"
|
||||
|
||||
namespace mujoco::studio {
|
||||
|
||||
class ObjectLauncher {
|
||||
public:
|
||||
ObjectLauncher() : rng_(std::random_device{}()) {}
|
||||
|
||||
void UpdateGui() {
|
||||
using platform::ImGui_Input;
|
||||
|
||||
ImGui::Checkbox("Enable Key Binding (Ctrl+Shift+Enter)", &enabled_);
|
||||
ImGui_Input("Size", &size_, {0.01f, 1.0f, 0.01, 0.1});
|
||||
ImGui_Input("Speed", &speed_, {0.01f, 100.0f, 0.1, 1.0});
|
||||
ImGui_Input("Mass", &mass_, {0.01f, 100.0f, 0.01, 0.1});
|
||||
ImGui_Input("Life", &lifetime_, {0.0f, 60.0f, 0.1, 1.0});
|
||||
|
||||
int shape = type_ == mjGEOM_BOX ? 0 : 1;
|
||||
const char* names[] = {"Box", "Sphere"};
|
||||
ImGui::Combo("Shape", &shape, names, 2);
|
||||
type_ = shape == 0 ? mjGEOM_BOX : mjGEOM_SPHERE;
|
||||
|
||||
if (ImGui::Button("Launch", ImVec2(-1.0f, 0.0f))) {
|
||||
active_ = true;
|
||||
}
|
||||
|
||||
if (ImGui::Button("Clear")) {
|
||||
for (auto& object : objects_) {
|
||||
object.expiration = -1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
void HandleKeyboardEvent() { if (enabled_) active_ = true; }
|
||||
|
||||
bool UpdateSpecPreCompile(mjSpec* spec, const mjModel* model,
|
||||
const mjData* data, const mjvCamera* camera) {
|
||||
// Remove expired objects.
|
||||
auto it = std::remove_if(
|
||||
objects_.begin(), objects_.end(), [&](const ObjectInfo& o) {
|
||||
const bool expired =
|
||||
o.body_id >= 0 && o.expiration != 0 && o.expiration < data->time;
|
||||
if (expired) {
|
||||
mjsBody* body = mjs_findBody(spec, o.name.c_str());
|
||||
if (body) {
|
||||
mjs_delete(spec, body->element);
|
||||
}
|
||||
}
|
||||
return expired;
|
||||
});
|
||||
if (it != objects_.end()) {
|
||||
objects_.erase(it, objects_.end());
|
||||
return true;
|
||||
}
|
||||
|
||||
if (!active_) return false;
|
||||
active_ = false;
|
||||
|
||||
mjsBody* world = mjs_findBody(spec, "world");
|
||||
if (!world) return false;
|
||||
mjsBody* body = mjs_addBody(world, nullptr);
|
||||
if (!body) return false;
|
||||
mjsJoint* joint = mjs_addJoint(body, nullptr);
|
||||
if (!joint) return false;
|
||||
mjsGeom* geom = mjs_addGeom(body, nullptr);
|
||||
if (!geom) return false;
|
||||
|
||||
ObjectInfo& object = objects_.emplace_back();
|
||||
object.name = "projectile" + std::to_string(counter_++);;
|
||||
object.expiration = data->time + lifetime_;
|
||||
|
||||
mjtNum pos[3];
|
||||
mjtNum dir[3];
|
||||
mjtNum up[3];
|
||||
mjv_cameraFrame(pos, dir, up, nullptr, data, camera);
|
||||
mjs_setName(body->element, object.name.c_str());
|
||||
|
||||
joint->type = mjJNT_FREE;
|
||||
body->mass = mass_;
|
||||
geom->type = type_;
|
||||
geom->size[0] = size_;
|
||||
geom->size[1] = size_;
|
||||
geom->size[2] = size_;
|
||||
// Slightly in front of the camera.
|
||||
body->pos[0] = pos[0] + (dir[0] * 0.1);
|
||||
body->pos[1] = pos[1] + (dir[1] * 0.1);
|
||||
body->pos[2] = pos[2] + (dir[2] * 0.1);
|
||||
// Randomize the color.
|
||||
geom->rgba[0] = std::uniform_real_distribution<float>(0.3f, 1.0f)(rng_);
|
||||
geom->rgba[1] = std::uniform_real_distribution<float>(0.3f, 1.0f)(rng_);
|
||||
geom->rgba[2] = std::uniform_real_distribution<float>(0.3f, 1.0f)(rng_);
|
||||
geom->rgba[3] = 1.0;
|
||||
// Launch it slightly upwards to get a nice arc.
|
||||
launch_vel_[0] = (dir[0] * speed_) + up[0];
|
||||
launch_vel_[1] = (dir[1] * speed_) + up[1];
|
||||
launch_vel_[2] = (dir[2] * speed_) + up[2];
|
||||
return true;
|
||||
}
|
||||
|
||||
void UpdateSpecPostCompile(const mjSpec* spec, const mjModel* model,
|
||||
mjData* data) {
|
||||
if (objects_.empty()) {
|
||||
return;
|
||||
}
|
||||
|
||||
ObjectInfo& object = objects_.back();
|
||||
if (object.launched) {
|
||||
return;
|
||||
}
|
||||
object.launched = true;
|
||||
|
||||
const int body_id = mj_name2id(model, mjOBJ_BODY, object.name.c_str());
|
||||
if (body_id < 0) {
|
||||
return;
|
||||
}
|
||||
int joint_id = model->body_jntadr[body_id];
|
||||
if (joint_id < 0 || model->jnt_type[joint_id] != mjJNT_FREE) {
|
||||
return;
|
||||
}
|
||||
int qvel_addr = model->jnt_dofadr[joint_id];
|
||||
if (qvel_addr < 0) {
|
||||
return;
|
||||
}
|
||||
object.body_id = body_id;
|
||||
data->qvel[qvel_addr + 0] = launch_vel_[0];
|
||||
data->qvel[qvel_addr + 1] = launch_vel_[1];
|
||||
data->qvel[qvel_addr + 2] = launch_vel_[2];
|
||||
}
|
||||
|
||||
private:
|
||||
struct ObjectInfo {
|
||||
std::string name;
|
||||
int body_id = -1;
|
||||
mjtNum expiration = 0;
|
||||
bool launched = false;
|
||||
};
|
||||
|
||||
std::mt19937 rng_;
|
||||
int counter_ = 0;
|
||||
bool enabled_ = false;
|
||||
bool active_ = false;
|
||||
mjtNum size_ = 0.13365;
|
||||
mjtNum speed_ = 10.0;
|
||||
mjtNum mass_ = 10.0;
|
||||
mjtNum lifetime_ = 5.0;
|
||||
mjtGeom type_ = mjGEOM_BOX;
|
||||
mjtNum launch_vel_[3] = {0, 0, 0};
|
||||
std::vector<ObjectInfo> objects_;
|
||||
};
|
||||
|
||||
} // namespace mujoco::studio
|
||||
|
||||
mjPLUGIN_LIB_INIT {
|
||||
using mujoco::studio::ObjectLauncher;
|
||||
|
||||
static ObjectLauncher plugin;
|
||||
|
||||
mujoco::platform::GuiPlugin gui;
|
||||
gui.data = &plugin;
|
||||
gui.name = "ObjectLauncher";
|
||||
gui.update = [](mujoco::platform::GuiPlugin* self) {
|
||||
auto* plugin = static_cast<mujoco::studio::ObjectLauncher*>(self->data);
|
||||
plugin->UpdateGui();
|
||||
};
|
||||
mujoco::platform::RegisterPlugin(gui);
|
||||
|
||||
mujoco::platform::KeyHandlerPlugin key_handler;
|
||||
key_handler.data = &plugin;
|
||||
key_handler.name = "ObjectLauncher";
|
||||
key_handler.key_chord = ImGuiKey_Enter | ImGuiMod_Ctrl | ImGuiMod_Shift;
|
||||
key_handler.on_key_pressed = [](mujoco::platform::KeyHandlerPlugin* self) {
|
||||
auto* plugin = static_cast<mujoco::studio::ObjectLauncher*>(self->data);
|
||||
plugin->HandleKeyboardEvent();
|
||||
};
|
||||
mujoco::platform::RegisterPlugin(key_handler);
|
||||
|
||||
mujoco::platform::SpecEditorPlugin spec_editor;
|
||||
spec_editor.data = &plugin;
|
||||
spec_editor.name = "ObjectLauncher";
|
||||
spec_editor.pre_compile = [](mujoco::platform::SpecEditorPlugin* self,
|
||||
mjSpec* spec, const mjModel* model,
|
||||
const mjData* data, const mjvCamera* camera) {
|
||||
auto* plugin = static_cast<mujoco::studio::ObjectLauncher*>(self->data);
|
||||
return plugin->UpdateSpecPreCompile(spec, model, data, camera);
|
||||
};
|
||||
spec_editor.post_compile = [](mujoco::platform::SpecEditorPlugin* self,
|
||||
const mjSpec* spec, const mjModel* model,
|
||||
mjData* data) {
|
||||
auto* plugin = static_cast<mujoco::studio::ObjectLauncher*>(self->data);
|
||||
return plugin->UpdateSpecPostCompile(spec, model, data);
|
||||
};
|
||||
mujoco::platform::RegisterPlugin(spec_editor);
|
||||
}
|
||||
Reference in New Issue
Block a user