Incorrect interaction between objects

[DINHQuangDung1999]

Hi everyone,

I am using a simple environment similar to the one in example_bounce.py. When I simulate the environment with obstacle that is not in the way of the ball, it seems that the ball still hit some invisible object.
Can you help me with this problem? I attached my code below

Bounce environment

import numpy as np
import os 
import warp as wp
import warp.sim
import warp.sim.render
from tqdm import trange

@wp.kernel
def loss_kernel(
    pos: wp.array(dtype=wp.vec3),
    target: wp.vec3,
    loss: wp.array(dtype=float),
):
    i = wp.tid()  # gets current thread id
    delta = pos[i] - target
    loss[i] = wp.dot(delta, delta)
    
@wp.kernel
def sum_kernel(
    losses: wp.array(dtype=float),
    loss: wp.array(dtype=float),
):
    i = wp.tid()
    wp.atomic_add(loss, 0, losses[i])

@wp.kernel
def mean_kernel(
    x_in: wp.array(dtype=wp.vec3), x_out: wp.array(dtype=wp.vec3), num_envs: float
):
    i = wp.tid()
    wp.atomic_add(x_out, 0, x_in[i] / num_envs)


@wp.kernel
def step_kernel(x: wp.array(dtype=wp.vec3), grad: wp.array(dtype=wp.vec3), alpha: float):
    tid = wp.tid()

    # gradient descent step
    x[tid] = x[tid] - grad[tid] * alpha

@wp.kernel
def noise_kernel(x: wp.array(dtype=wp.vec3), noise: wp.array(dtype=wp.vec3)):
    tid = wp.tid()
    x[tid] = x[tid] + noise[tid] 

@wp.kernel
def repeat_kernel(x: wp.array(dtype=wp.vec3f),
                  out: wp.array(dtype=wp.vec3f)):
    tid = wp.tid()
    out[tid] = x[0]

class Bounce:
    def __init__(self, 
                render = False, 
                num_envs = 128, 
                num_steps = 200,
                verbose=False, 
                zero_order = False,
                start_state=np.array([-0.5, 1.0, 0.0]),
                start_vel=np.array([5.0, -5.0, 0.0]),
                profile=False,
                soft_contact_ke=1e4,  # stiffness
                soft_contact_kf=0.0,  # stiffness of friction
                soft_contact_kd=1e1,  # damping
                soft_contact_mu=0.2,  # friction coefficient
                soft_contact_margin=1e1,
                std=1e-1,
                wall = None,
                target = (-2.0, 1.5, 0.0)):
        
        # flags 
        self.profile = profile
        self.verbose = verbose
        self.zero_order = zero_order
        self.render_trajectory = render 

        # time
        fps = 60
        self.frame_dt = 1.0 / fps
        frame_steps = num_steps

        self.sim_substeps = 8
        self.sim_steps = frame_steps * self.sim_substeps
        self.sim_dt = self.frame_dt / self.sim_substeps

        self.iter = 0
        self.render_time = 0.0

        self.train_rate = 0.02

        # scene
        self.num_envs = num_envs
        np.random.seed(0)
        self.std = std
        noise = self.noise()
        builder = wp.sim.ModelBuilder()
        for i in range(num_envs):
            builder.add_particle(pos=start_state, vel=start_vel + noise[i], mass=1.0)

        self.wall = wall
        if self.wall == 'high':
            # high wall
            builder.add_shape_box(body=-1, pos=wp.vec3(2.0, 0.65, 2.0), hx=0.25, hy=0.65, hz=0.5)
        elif self.wall == 'low':
            # short wall
            builder.add_shape_box(body=-1, pos=(2.0, 0.5, 0.0), hx=0.25, hy=0.5, hz=0.5)
        self.target = target
        self.start_vel = start_vel
        self.start_state = start_state

        # use `requires_grad=True` to create a model for differentiable simulation
        self.model = builder.finalize(requires_grad=True)
        self.model.ground = True

        self.model.soft_contact_ke = soft_contact_ke
        self.model.soft_contact_kf = soft_contact_kf
        self.model.soft_contact_kd = soft_contact_kd
        self.model.soft_contact_mu = soft_contact_mu
        self.model.soft_contact_margin = soft_contact_margin
        self.model.soft_contact_restitution = 1.0

        self.integrator = wp.sim.SemiImplicitIntegrator()

        # Allocate loss and mean grad
        self.loss = wp.zeros(self.num_envs, dtype=wp.float32, requires_grad=True)
        self.l = wp.zeros(1, dtype=wp.float32, requires_grad=True)
        self.m_grad = wp.zeros(1, dtype=wp.vec3f)
        self.mean_grad = wp.zeros(self.num_envs, dtype=wp.vec3f)

        # allocate sim states for trajectory
        self.states = []
        for _i in range(self.sim_steps + 1):
            self.states.append(self.model.state())

        # record
        self.trajectories = []
        self.grad_norms = []
        self.loss_history = []

        # one-shot contact creation (valid if we're doing simple collision against a constant normal plane)
        wp.sim.collide(self.model, self.states[0])

        if self.render_trajectory:
            self.renderer = wp.sim.render.SimRenderer(self.model,
                    os.path.join(
                                    os.path.dirname(__file__), "outputs/example_sim_grad_bounce.usd"
                                ), 
                scaling=1.0)
        else:
            self.renderer = None

        # capture forward/backward passes
        self.use_cuda_graph = True
        if self.use_cuda_graph:
            with wp.ScopedCapture() as capture:
                self.tape = wp.Tape()
                with self.tape:
                    self.forward()
                if not self.zero_order:
                    self.tape.backward(self.l)
            self.graph = capture.graph

    def noise(self):
        noise = np.random.normal(0.0, self.std, (self.num_envs, 2))
        noise[0] = 0.0  # for baseline
        self.noise_ = np.append(noise, np.zeros((self.num_envs, 1)), axis=1)
        return self.noise_

    def reset(self, start_state=None, start_vel=None):
        if start_state is None:
            start_state = self.start_state

        if start_vel is None:
            start_vel = self.start_vel

        # replicate array if necessary and assign
        q = np.array(start_state)
        if q.ndim == 1:
            q = np.tile(q, (self.num_envs, 1))
        q = wp.array(q, dtype=wp.vec3)
        self.model.particle_q.assign(q)

        # replicate array if necessary and assign
        qd = np.array(start_vel)
        if qd.ndim == 1:
            qd = np.tile(qd, (self.num_envs, 1))
        qd += self.noise()
        qd = wp.array(qd, dtype=wp.vec3)
        self.model.particle_qd.assign(qd)

        # only need to reset first state as everything is forward simulated from it
        self.states[0] = self.model.state(requires_grad=True)

        self.loss.zero_()
        self.l.zero_()

    @wp.kernel
    def count_contact_changes_kernel(
        contact_count: wp.array(dtype=int),
        contact_count_id: int,
        contact_count_copy: wp.array(dtype=int),
    ):
        wp.atomic_add(contact_count_copy, contact_count_id, contact_count[0])

    def count_contact_changes(self, i):
        # count contact changes
        wp.launch(
            self.count_contact_changes_kernel,
            dim=self.num_envs,
            inputs=[self.model.soft_contact_count, i],
            outputs=[self.contact_count],
            device=self.device,
        )

    def compute_loss(self):
        # run control loop
        for i in range(self.sim_steps):
            self.states[i].clear_forces()
            self.integrator.simulate(self.model, self.states[i], self.states[i + 1], self.sim_dt)

        # compute loss on final state
        wp.launch(loss_kernel, dim=self.num_envs, inputs=[self.states[-1].particle_q, self.target, self.loss])
        return self.loss

    def sum_loss(self):
        wp.launch(sum_kernel, dim=self.num_envs, inputs=[self.loss, self.l])
        return self.l 
    
    def forward(self):
        self.compute_loss()
        self.sum_loss()

    def get_trajectory(self):
        xy = []
        for s in self.states:
            pos = s.particle_q.numpy()[:, :2]
            vel = s.particle_qd.numpy()[:, :2]
            xy.append(np.concatenate((pos, vel), axis=-1))
        return np.array(xy)
    
    def run_env(self):
        # run control loop
        for i in range(self.sim_steps):
            self.states[i].clear_forces()
            self.integrator.simulate(self.model, self.states[i], self.states[i + 1], self.sim_dt)

    def train(self, iters):
        vel_history = []
        for i in trange(iters):
            vel_history.append(self.states[0].particle_qd.numpy())
            with wp.ScopedTimer("step", active=self.verbose):
                if self.use_cuda_graph:
                    wp.capture_launch(self.graph)
                else:
                    self.tape = wp.Tape()
                    with self.tape:
                        self.forward()
                    if self.zero_order:
                        self.tape.backward(self.l)    

            x = self.states[0].particle_qd
            # gradient descent step
            if self.zero_order:
                loss = self.loss.numpy()
                baseline = loss[0]
                x_grad = (
                    1/self.std**2
                    * (loss[..., None] - baseline)
                    * self.noise_
                )
                x_grad = x_grad.mean(axis=0)
                self.grad_norms.append(np.linalg.norm(x_grad))
                x_grad = wp.from_numpy(np.repeat([x_grad], self.num_envs), dtype=wp.vec3f)
                wp.launch(step_kernel, dim=len(x), inputs=[x, x_grad, self.train_rate])
            else:
                wp.launch(mean_kernel, dim=len(x), inputs = [x.grad, self.m_grad, self.num_envs])
                wp.launch(repeat_kernel, dim=self.num_envs, inputs = [self.m_grad, self.mean_grad])
                wp.launch(step_kernel, dim=len(x), inputs=[x, self.mean_grad, self.train_rate])
                
                self.grad_norms.append(np.linalg.norm(self.m_grad.numpy()))
                self.tape.zero() # clear grads for next iteration
                self.m_grad.zero_()
                self.mean_grad.zero_()
                self.l.zero_()

            total_loss = self.loss.numpy()
            if self.verbose:
                print(f"Iter: {self.iter} Loss: {total_loss.mean()}")
                #     print(f"    x: {x} g: {x_grad}")

            # add noise to reset
            new_noise = wp.from_numpy(self.noise(), dtype=wp.vec3f)
            wp.launch(noise_kernel, dim = len(x), inputs = [x, new_noise])

            self.loss_history.append(total_loss)
            self.trajectories.append(self.get_trajectory())
            self.iter = self.iter + 1

        mean_loss = np.array(self.loss_history).mean(axis = 1)
        min_mean_loss = self.loss_history[np.argmin(mean_loss)]
        min_loss_idx = np.argmin(min_mean_loss)
        self.best_v0 = vel_history[np.argmin(mean_loss)][min_loss_idx]
        print('Best initial velocity: {} at iter {}. Loss: {}.'.format(self.best_v0, 
                                                                       np.argmin(mean_loss),
                                                                       self.loss_history[np.argmin(mean_loss)][min_loss_idx]))
        return np.array(self.loss_history), np.array(self.trajectories), np.array(self.grad_norms)
    
    def render(self, env_id = 0):
        if self.renderer is None:
            return

        with wp.ScopedTimer("render"):
            # draw trajectory (of only the environment with no noise, i.e. 0)
            traj_verts = [self.states[0].particle_q.numpy()[env_id].tolist()]

            for i in range(0, self.sim_steps, self.sim_substeps):
                traj_verts.append(self.states[i].particle_q.numpy()[env_id].tolist())

                self.renderer.begin_frame(self.render_time)
                state_i = wp.sim.State()
                q_state_i = self.states[i].particle_q.numpy()[env_id]
                state_i.particle_q = wp.from_numpy(q_state_i, dtype=wp.vec3f)
                self.renderer.render(state_i)
                # self.renderer.render(self.states[i])
                self.renderer.render_box(
                    pos=self.target,
                    rot=wp.quat_identity(),
                    extents=(0.1, 0.1, 0.1),
                    name="target",
                    color=(0.0, 0.0, 0.0),
                )
                self.renderer.render_line_strip(
                    vertices=traj_verts,
                    color=wp.render.bourke_color_map(0.0, 7.0, self.loss.numpy()[0]),
                    radius=0.02,
                    name=f"traj_{self.iter}",
                )
                self.renderer.end_frame()

                from pxr import Gf, UsdGeom

                particles_prim = self.renderer.stage.GetPrimAtPath("/root/particles")
                particles = UsdGeom.Points.Get(self.renderer.stage, particles_prim.GetPath())
                particles.CreateDisplayColorAttr().Set([Gf.Vec3f(1.0, 1.0, 1.0)], time=self.renderer.time)

                self.render_time += self.frame_dt

    def check_grad(self):
        param = self.states[0].particle_qd

        # initial value
        x_c = param.numpy().flatten()

        # compute numeric gradient
        x_grad_numeric = np.zeros_like(x_c)

        for i in range(len(x_c)):
            eps = 1.0e-3

            step = np.zeros_like(x_c)
            step[i] = eps

            x_1 = x_c + step
            x_0 = x_c - step

            param.assign(x_1)
            l_1 = self.forward().numpy()[0]

            param.assign(x_0)
            l_0 = self.forward().numpy()[0]

            dldx = (l_1 - l_0) / (eps * 2.0)

            x_grad_numeric[i] = dldx

        # reset initial state
        param.assign(x_c)

        # compute analytic gradient
        tape = wp.Tape()
        with tape:
            l = self.forward()

        tape.backward(l)

        x_grad_analytic = tape.gradients[param]

        print(f"numeric grad: {x_grad_numeric}")
        print(f"analytic grad: {x_grad_analytic}")

        tape.zero()

Code for running the environment

# Copyright (c) 2022 NVIDIA CORPORATION.  All rights reserved.
# NVIDIA CORPORATION and its licensors retain all intellectual property
# and proprietary rights in and to this software, related documentation
# and any modifications thereto.  Any use, reproduction, disclosure or
# distribution of this software and related documentation without an express
# license agreement from NVIDIA CORPORATION is strictly prohibited.

import os
import numpy as np
from tqdm import tqdm
import torch
from time import time

from bounce_env_narrow import Bounce

import warp as wp
import warp.sim
import warp.sim.render

def main(cfg):
    np.random.seed(0)

    print("First-order optimisation")
    env = Bounce(num_envs=1, num_steps=80, std=0.1, wall = 'high', render = True, 
                    start_vel=np.array([ 9.759769  ,-4.2871647  ,0.       ]))
    env.run_env()
    env.render()
    env.renderer.save()
        

if __name__ == "__main__":

    main()

[DINHQuangDung1999]

Bonus: Here is the rendered trajectory when I create two walls.

[DINHQuangDung1999]

Last update: If I try running example_bounce.py where the wall position is changed to (2.0, 0.65, 3.0), the problem still exists. I wonder what is the reason behind this.