diff --git a/kubric/simulator/pybullet.py b/kubric/simulator/pybullet.py index 58798597..3233a1ef 100644 --- a/kubric/simulator/pybullet.py +++ b/kubric/simulator/pybullet.py @@ -20,6 +20,7 @@ import tempfile from typing import Dict, List, Optional, Tuple, Union +import numpy as np import tensorflow as tf from singledispatchmethod import singledispatchmethod @@ -175,7 +176,8 @@ def save_state(self, path: Union[pathlib.Path, str] = "scene.bullet"): def run( self, frame_start: int = 0, - frame_end: Optional[int] = None + frame_end: Optional[int] = None, + forces: Dict[int, Dict[int, Tuple[Tuple[float, float, float], float]]] = {}, ) -> Tuple[Dict[core.PhysicalObject, Dict[str, list]], List[dict]]: """ Run the physics simulation. @@ -227,8 +229,22 @@ def run( }) if current_step % steps_per_frame == 0: + current_frame = current_step // steps_per_frame + for obj_idx in obj_idxs: position, quaternion = self.get_position_and_rotation(obj_idx) + if obj_idx in forces and current_frame in forces[obj_idx]: + force, bias = forces[obj_idx][current_frame] + # pb.applyExternalForce(objectUniqueId=obj_idx, + # linkIndex=-1, + # forceObj=force, + # posObj=position, + # flags=pb.WORLD_FRAME) + effective_force = np.array(force) - bias * np.array(position) + pb.resetBaseVelocity(objectUniqueId=obj_idx, + linearVelocity=effective_force) + + velocity, angular_velocity = self.get_velocities(obj_idx) animation[obj_idx]["position"].append(position)